From 2a40260f8ebc518638b3fa805bfac1a5a117528b Mon Sep 17 00:00:00 2001 From: Reiase Date: Thu, 20 Aug 2026 08:37:17 +0800 Subject: [PATCH 1/7] refactor(pchronicle): unify unknown fields and JSON streaming --- crates/persisting-pchronicle/Cargo.toml | 4 + crates/persisting-pchronicle/README.md | 36 +- .../benches/json_streaming.rs | 368 ++++- .../src/agenticmd/codec.rs | 38 +- .../src/agenticmd/convert.rs | 365 +++-- .../src/agenticmd/validate.rs | 50 +- crates/persisting-pchronicle/src/atif.rs | 118 +- .../persisting-pchronicle/src/convert/actf.rs | 752 ++++++----- .../persisting-pchronicle/src/convert/atif.rs | 638 +++++---- .../src/convert/events.rs | 3 +- crates/persisting-pchronicle/src/document.rs | 284 ++-- .../persisting-pchronicle/src/formats/mod.rs | 9 +- .../src/formats/openai_corpus.rs | 1188 +++++++++-------- .../src/formats/storyline.rs | 316 ++--- .../src/formats/unknown_fields.rs | 1081 +++++++++++++++ crates/persisting-pchronicle/src/lib.rs | 8 +- crates/persisting-pchronicle/src/model.rs | 9 +- .../src/store/catalog/tests.rs | 3 +- .../src/store/document_source.rs | 26 +- .../src/store/files/actf_reader.rs | 176 +++ .../src/store/files/actf_stream.rs | 868 ++++++++++++ .../src/store/files/atif_reader.rs | 379 ++++-- .../src/store/files/atif_stream.rs | 823 ++++-------- .../src/store/files/json_stream.rs | 609 +++++++++ .../src/store/files/mod.rs | 164 ++- .../src/store/files/projected_steps.rs | 204 +++ .../src/store/files/tests.rs | 114 ++ .../src/store/storyline/content.rs | 203 ++- .../src/store/storyline/mod.rs | 30 +- .../src/store/storyline/model.rs | 76 +- .../src/store/storyline/mutation.rs | 8 +- .../src/store/storyline/rows.rs | 57 +- .../src/store/storyline/tests.rs | 200 ++- crates/persisting-pchronicle/src/tests.rs | 85 +- crates/persisting-pchronicle/tests/README.md | 2 + .../tests/atif_lance_corpus.rs | 10 +- .../tests/conversion_semantics.rs | 421 ++++++ .../tests/document_source.rs | 45 +- .../tests/import_roundtrip_fixtures.rs | 6 +- .../tests/query_engine.rs | 6 + .../tests/storyline_lance_roundtrip.rs | 13 +- .../tests/unknown_fields_roundtrip.rs | 58 + docs/src/pchronicle/design/storyline-lance.md | 17 +- .../pchronicle/design/storyline-lance.zh.md | 17 +- .../pchronicle/design/trajectory-storage.md | 4 +- .../design/trajectory-storage.zh.md | 4 +- docs/src/rfcs/0004-actf-format.md | 19 +- 47 files changed, 7247 insertions(+), 2667 deletions(-) create mode 100644 crates/persisting-pchronicle/src/formats/unknown_fields.rs create mode 100644 crates/persisting-pchronicle/src/store/files/actf_reader.rs create mode 100644 crates/persisting-pchronicle/src/store/files/actf_stream.rs create mode 100644 crates/persisting-pchronicle/src/store/files/json_stream.rs create mode 100644 crates/persisting-pchronicle/src/store/files/projected_steps.rs create mode 100644 crates/persisting-pchronicle/tests/conversion_semantics.rs create mode 100644 crates/persisting-pchronicle/tests/unknown_fields_roundtrip.rs diff --git a/crates/persisting-pchronicle/Cargo.toml b/crates/persisting-pchronicle/Cargo.toml index dc57e118..8b87c570 100644 --- a/crates/persisting-pchronicle/Cargo.toml +++ b/crates/persisting-pchronicle/Cargo.toml @@ -92,6 +92,10 @@ required-features = ["search"] name = "storyline_lance_roundtrip" required-features = ["lance-store"] +[[test]] +name = "conversion_semantics" +required-features = ["lance-store"] + [[test]] name = "atif_lance_corpus" required-features = ["lance-store"] diff --git a/crates/persisting-pchronicle/README.md b/crates/persisting-pchronicle/README.md index 527afbb8..983093e8 100644 --- a/crates/persisting-pchronicle/README.md +++ b/crates/persisting-pchronicle/README.md @@ -31,8 +31,8 @@ events.lance ──单向投影──► StorylineDocument ──► Storyline | `Storyline` | Storyline 三表 Lance | `runs`、`steps`、`tool_calls` | 权威二进制表示 | | `AgenticMd` | Markdown 文件 | `runs`、`steps`、`tool_calls` | 可读编码，可双向转换 | | `Atif` | ATIF JSON/JSONL/NDJSON | `runs`、`steps`、`tool_calls` | ATIF v1.7 对齐，可双向转换 | -| `OpenaiMsg` | OpenAI message corpus JSON | `runs`、`steps`、`tool_calls` | 通过分层 residual 无损往返 | -| `Actf` | ACTF JSON | `runs`、`steps`、`tool_calls` | 通过分层 residual 无损往返 | +| `OpenaiMsg` | OpenAI message corpus JSON | `runs`、`steps`、`tool_calls` | 通过分层 unknown fields 无损往返 | +| `Actf` | ACTF JSON | `runs`、`steps`、`tool_calls` | 通过分层 unknown fields 无损往返 | 统一读取入口是 `document::open_document`；返回的 `DocumentSource` 隐藏具体 provider， 并提供有预算上限的物化、逐条 Storyline 回调和 DataFusion 注册。写入仍使用 @@ -48,17 +48,27 @@ ACTF → Storyline Lance → ACTF OpenAI Msg → Storyline Lance → OpenAI Msg ``` -保真内容包括 ATIF 的 missing/null/value 三态、嵌套 subagent 顺序、`trajectory_id` 与 -run-scoped `session_id` 的独立身份、RFC3339 原始偏移与亚毫秒精度，以及 ACTF/OpenAI 的 -未知字段、数组顺序、attempt 分组和多 session 关系。外围格式无法映射到正式 Storyline -字段的内容保存在对应语义层级的受控 residual 中；不会保存完整原始对象副本。跨格式转换 -只保证目标格式能够表达的语义。Canonical Event → Storyline 是有意的有损规范化投影， -不属于上述无损承诺。 +已建模的 Storyline 语义（包括嵌套 subagent 顺序、`trajectory_id` 与 run-scoped +`session_id` 的独立身份、RFC3339 原始偏移与亚毫秒精度，以及 ACTF/OpenAI 的数组顺序、 +attempt 分组和多 session 关系）按其规范化表示保存。已知字段的 missing/null 区别，以及 +输入的物理容器形态（例如 ATIF 顶层单对象与单元素数组），都会被规范化，因而不作为 +往返保真承诺。 -ATIF 的顶层单对象/数组形态与 root 顺序作为格式无关的 Storyline 集合语义随 Lance -持久化，不依赖进程内 sidecar；Lance 内部另用 `storage_ordinal` 维护全局稳定读取顺序， -不会把多次增量写入都退化到 document id 排序。无法用任何 Storyline 表达的空 ATIF 数组 -或空 OpenAI 信封会 fail closed，而不是接受后在导出时静默丢失容器字段。 +源格式中 Storyline 未建模的键保存在受控 unknown fields：键名是带命名空间的精确 +[RFC 6901 JSON Pointer](https://www.rfc-editor.org/rfc/rfc6901)；未知字段值不保存完整原始对象 +副本。未知键即使值为 `null` 也会保留。写回同一格式时，目标格式的规范字段优先；若 +unknown field 与它们冲突，编码会 fail closed，而不会覆盖目标字段或静默丢弃冲突。 + +跨格式、多跳转换使用保留的 version-1 `_storyline` envelope 携带这些 unknown fields，确保目标 +格式不能直接表示的源语义仍可在后续转换中恢复。每条 trajectory 跨所有来源默认最多 +4,096 个 unknown fields、最多 1 MiB；任一上限溢出都会拒绝整条 Storyline，而非截断或只 +保留部分未知字段。Canonical Event → Storyline 是有意的有损规范化投影，不属于上述无损 +承诺。 + +Storyline Lance 的 `objects.lance` 可用于 unknown field 值的内部去重/卸载优化；它从不出现在 +公共 Storyline 模型或任何公共 wire 输出中。Lance 内部另用 `storage_ordinal` 维护全局稳定 +读取顺序，不会把多次增量写入都退化到 document id 排序。无法用任何 Storyline 表达的空 +ATIF 数组或空 OpenAI 信封会 fail closed，而不是接受后在导出时静默丢失容器字段。 ## DataFusion 能力 @@ -70,7 +80,7 @@ ATIF 的顶层单对象/数组形态与 root 顺序作为格式无关的 Storyli | Storyline Lance | 是 | expression-dependent | 是 | 是 | 否 | 是 | | ATIF | 是 | inexact | 是 | 否 | 是 | 否 | | OpenAI Msg | 是 | unsupported | 是 | 否 | 否 | 否 | -| ACTF | 是 | unsupported | 是 | 否 | 否 | 否 | +| ACTF | 是 | inexact | 是 | 否 | 是 | 否 | | AgenticMD | 是 | unsupported | 否 | 否 | 否 | 否 | Canonical Event 保留 Lance projection/filter/limit pushdown、scalar index、pinned manifest diff --git a/crates/persisting-pchronicle/benches/json_streaming.rs b/crates/persisting-pchronicle/benches/json_streaming.rs index 8d239482..2c1c05f9 100644 --- a/crates/persisting-pchronicle/benches/json_streaming.rs +++ b/crates/persisting-pchronicle/benches/json_streaming.rs @@ -1,5 +1,12 @@ //! Allocation, throughput, tail-latency, and process-RSS benchmark for a cold -//! projected ATIF JSON datasource + DataFusion query iteration. +//! projected JSON datasource + DataFusion query iteration. +//! +//! Environment variables: +//! - `PCHRONICLE_BENCH_SCALE` (default: 128) +//! - `PCHRONICLE_BENCH_ITERS` (default: 20) +//! - `PCHRONICLE_BENCH_JSON_SHAPE`: `ndjson` | `array` (default: `ndjson`) +//! - `PCHRONICLE_BENCH_FORMAT`: `atif` | `actf` (default: `atif`) +//! - `PCHRONICLE_BENCH_PATH`: `projected` | `full` | `both` (default: `projected`) use std::alloc::{GlobalAlloc, Layout, System}; use std::fs::File; @@ -55,13 +62,46 @@ unsafe impl GlobalAlloc for CountingAllocator { #[global_allocator] static GLOBAL_ALLOCATOR: CountingAllocator = CountingAllocator; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum BenchFormat { + Atif, + Actf, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum BenchPath { + Projected, + Full, +} + #[derive(Debug)] struct Sample { elapsed: Duration, allocations: u64, allocated_bytes: u64, rows_scanned: u64, + files_parsed: u64, + input_buffer_peak_bytes: u64, + projected_files: u64, +} + +#[derive(Debug)] +struct BenchReport { + format: BenchFormat, + path: BenchPath, + json_shape: String, + documents: usize, + steps: u64, + file_bytes: u64, + iterations: usize, + median: Duration, + p95: Duration, + median_allocations: u64, + p95_allocations: u64, + median_allocated_bytes: u64, + p95_allocated_bytes: u64, input_buffer_peak_bytes: u64, + rss_mib: f64, } fn main() -> Result<()> { @@ -69,6 +109,12 @@ fn main() -> Result<()> { let iterations = env_usize("PCHRONICLE_BENCH_ITERS", 20); let json_shape = std::env::var("PCHRONICLE_BENCH_JSON_SHAPE").unwrap_or_else(|_| "ndjson".to_string()); + let format = parse_format( + &std::env::var("PCHRONICLE_BENCH_FORMAT").unwrap_or_else(|_| "atif".to_string()), + )?; + let path_mode = parse_path_mode( + &std::env::var("PCHRONICLE_BENCH_PATH").unwrap_or_else(|_| "projected".to_string()), + )?; anyhow::ensure!(scale > 0, "PCHRONICLE_BENCH_SCALE must be positive"); anyhow::ensure!(iterations > 0, "PCHRONICLE_BENCH_ITERS must be positive"); anyhow::ensure!( @@ -82,86 +128,264 @@ fn main() -> Result<()> { } else { "streaming.ndjson" }); - let (documents, steps) = write_corpus(&input, scale, &json_shape)?; + let (documents, steps) = write_corpus(&input, scale, &json_shape, format)?; + let file_bytes = std::fs::metadata(&input)?.len(); let runtime = tokio::runtime::Builder::new_multi_thread() .enable_all() .build()?; - runtime.block_on(run_query(&input))?; - let mut samples = Vec::with_capacity(iterations); - for _ in 0..iterations { + let paths = match path_mode { + BenchPathMode::Projected => vec![BenchPath::Projected], + BenchPathMode::Full => vec![BenchPath::Full], + BenchPathMode::Both => vec![BenchPath::Projected, BenchPath::Full], + }; + let dataset = BenchDataset { + json_shape: &json_shape, + input: &input, + documents, + steps, + file_bytes, + iterations, + }; + + let mut reports = Vec::with_capacity(paths.len()); + for path in paths { + runtime.block_on(run_query(format, path, &input))?; + let report = benchmark_path(&runtime, format, path, &dataset)?; + print_report(&report); + reports.push(report); + } + + if reports.len() == 2 { + print_comparison(&reports[0], &reports[1]); + } + Ok(()) +} + +struct BenchDataset<'a> { + json_shape: &'a str, + input: &'a std::path::Path, + documents: usize, + steps: u64, + file_bytes: u64, + iterations: usize, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum BenchPathMode { + Projected, + Full, + Both, +} + +fn parse_format(value: &str) -> Result { + match value { + "atif" => Ok(BenchFormat::Atif), + "actf" => Ok(BenchFormat::Actf), + other => anyhow::bail!("PCHRONICLE_BENCH_FORMAT must be 'atif' or 'actf', got '{other}'"), + } +} + +fn parse_path_mode(value: &str) -> Result { + match value { + "projected" => Ok(BenchPathMode::Projected), + "full" => Ok(BenchPathMode::Full), + "both" => Ok(BenchPathMode::Both), + other => anyhow::bail!( + "PCHRONICLE_BENCH_PATH must be 'projected', 'full', or 'both', got '{other}'" + ), + } +} + +fn benchmark_path( + runtime: &tokio::runtime::Runtime, + format: BenchFormat, + path: BenchPath, + dataset: &BenchDataset<'_>, +) -> Result { + let mut samples = Vec::with_capacity(dataset.iterations); + for _ in 0..dataset.iterations { reset_allocations(); let started = Instant::now(); - let query = runtime.block_on(run_query(&input))?; + let query = runtime.block_on(run_query(format, path, dataset.input))?; let elapsed = started.elapsed(); let (allocations, allocated_bytes) = allocation_snapshot(); samples.push(Sample { elapsed, allocations, allocated_bytes, - rows_scanned: query.0, - input_buffer_peak_bytes: query.1, + rows_scanned: query.rows_scanned, + files_parsed: query.files_parsed, + input_buffer_peak_bytes: query.streaming_buffer_peak_bytes, + projected_files: query.projected_files, }); } - anyhow::ensure!(samples.iter().all(|sample| sample.rows_scanned == steps)); + match path { + BenchPath::Projected => { + anyhow::ensure!( + samples.iter().all(|sample| sample.projected_files > 0), + "projected path did not record projected_files" + ); + anyhow::ensure!( + samples + .iter() + .all(|sample| sample.rows_scanned == dataset.steps), + "projected path rows_scanned mismatch" + ); + } + BenchPath::Full => { + anyhow::ensure!( + samples.iter().all(|sample| sample.projected_files == 0), + "full path unexpectedly used projected streaming" + ); + anyhow::ensure!( + samples.iter().all(|sample| sample.files_parsed > 0), + "full path did not parse any files" + ); + } + } + let median = percentile_duration(&samples, 0.50); let p95 = percentile_duration(&samples, 0.95); - let median_allocations = percentile_u64(&samples, 0.50, |sample| sample.allocations); - let p95_allocations = percentile_u64(&samples, 0.95, |sample| sample.allocations); - let median_allocated_bytes = percentile_u64(&samples, 0.50, |sample| sample.allocated_bytes); - let p95_allocated_bytes = percentile_u64(&samples, 0.95, |sample| sample.allocated_bytes); - let input_buffer_peak_bytes = samples - .iter() - .map(|sample| sample.input_buffer_peak_bytes) - .max() - .unwrap_or_default(); - let rows_per_second = steps as f64 / median.as_secs_f64(); - let rss_mib = process_peak_rss_bytes()? as f64 / (1024.0 * 1024.0); + Ok(BenchReport { + format, + path, + json_shape: dataset.json_shape.to_string(), + documents: dataset.documents, + steps: dataset.steps, + file_bytes: dataset.file_bytes, + iterations: dataset.iterations, + median, + p95, + median_allocations: percentile_u64(&samples, 0.50, |sample| sample.allocations), + p95_allocations: percentile_u64(&samples, 0.95, |sample| sample.allocations), + median_allocated_bytes: percentile_u64(&samples, 0.50, |sample| sample.allocated_bytes), + p95_allocated_bytes: percentile_u64(&samples, 0.95, |sample| sample.allocated_bytes), + input_buffer_peak_bytes: samples + .iter() + .map(|sample| sample.input_buffer_peak_bytes) + .max() + .unwrap_or_default(), + rss_mib: process_peak_rss_bytes()? as f64 / (1024.0 * 1024.0), + }) +} +fn print_report(report: &BenchReport) { + let rows_per_second = report.steps as f64 / report.median.as_secs_f64(); + let format_name = match report.format { + BenchFormat::Atif => "atif", + BenchFormat::Actf => "actf", + }; + let path_name = match report.path { + BenchPath::Projected => "projected", + BenchPath::Full => "full", + }; println!( - "dataset: shape={json_shape}, {documents} trajectories, {steps} steps, {} bytes", - std::fs::metadata(&input)?.len() + "dataset: format={format_name}, path={path_name}, shape={}, {} trajectories, {} steps, {} bytes", + report.json_shape, report.documents, report.steps, report.file_bytes ); println!( - "projected JSON: median={:.3} ms p95={:.3} ms rows/s={rows_per_second:.0}", - milliseconds(median), - milliseconds(p95) + "{path_name} JSON: median={:.3} ms p95={:.3} ms rows/s={rows_per_second:.0}", + milliseconds(report.median), + milliseconds(report.p95) ); println!( - "allocation traffic: median={median_allocations} calls/{median_allocated_bytes} bytes, \ - p95={p95_allocations} calls/{p95_allocated_bytes} bytes" + "allocation traffic: median={} calls/{} bytes, p95={} calls/{} bytes", + report.median_allocations, + report.median_allocated_bytes, + report.p95_allocations, + report.p95_allocated_bytes ); println!( - "memory: process peak RSS={rss_mib:.3} MiB, input buffer peak={input_buffer_peak_bytes} bytes" + "memory: process peak RSS={:.3} MiB, input buffer peak={} bytes", + report.rss_mib, report.input_buffer_peak_bytes ); println!( - "RESULT benchmark=json_streaming shape={json_shape} documents={documents} rows={steps} iterations={iterations} \ + "RESULT benchmark=json_streaming format={format_name} path={path_name} shape={} documents={} rows={} iterations={} \ median_ms={:.3} p95_ms={:.3} rows_s={rows_per_second:.0} \ - median_allocations={median_allocations} p95_allocations={p95_allocations} \ - median_allocated_bytes={median_allocated_bytes} p95_allocated_bytes={p95_allocated_bytes} \ - process_peak_rss_mib={rss_mib:.3} input_buffer_peak_bytes={input_buffer_peak_bytes}", - milliseconds(median), - milliseconds(p95), + median_allocations={} p95_allocations={} \ + median_allocated_bytes={} p95_allocated_bytes={} \ + process_peak_rss_mib={:.3} input_buffer_peak_bytes={}", + report.json_shape, + report.documents, + report.steps, + report.iterations, + milliseconds(report.median), + milliseconds(report.p95), + report.median_allocations, + report.p95_allocations, + report.median_allocated_bytes, + report.p95_allocated_bytes, + report.rss_mib, + report.input_buffer_peak_bytes + ); +} + +fn print_comparison(projected: &BenchReport, full: &BenchReport) { + let speedup = full.median.as_secs_f64() / projected.median.as_secs_f64(); + let alloc_ratio = full.median_allocated_bytes as f64 / projected.median_allocated_bytes as f64; + println!( + "comparison: projected median {:.3} ms vs full {:.3} ms ({speedup:.2}x faster), \ + allocated bytes ratio {alloc_ratio:.2}x (full/projected)", + milliseconds(projected.median), + milliseconds(full.median), ); - Ok(()) } -async fn run_query(input: &std::path::Path) -> Result<(u64, u64)> { +struct QueryMetrics { + rows_scanned: u64, + files_parsed: u64, + streaming_buffer_peak_bytes: u64, + projected_files: u64, +} + +async fn run_query( + format: BenchFormat, + path: BenchPath, + input: &std::path::Path, +) -> Result { + let document_format = match format { + BenchFormat::Atif => DocumentFormat::Atif, + BenchFormat::Actf => DocumentFormat::Actf, + }; let engine = ChronicleQueryEngine::open( - DocumentFormat::Atif, + document_format, input, ChronicleQueryExecutionOptions::default(), ) .await?; - engine - .query("SELECT source, COUNT(*) FROM steps GROUP BY source") - .await?; - let metrics = engine.local_file_metrics().expect("ATIF file metrics"); - Ok((metrics.rows_scanned, metrics.streaming_buffer_peak_bytes)) + let sql = match path { + BenchPath::Projected => "SELECT source, COUNT(*) FROM steps GROUP BY source", + BenchPath::Full => "SELECT * FROM steps", + }; + engine.query(sql).await?; + let metrics = engine.local_file_metrics().expect("local file metrics"); + Ok(QueryMetrics { + rows_scanned: metrics.rows_scanned, + files_parsed: metrics.files_parsed, + streaming_buffer_peak_bytes: metrics.streaming_buffer_peak_bytes, + projected_files: metrics.projected_files, + }) +} + +fn write_corpus( + path: &std::path::Path, + scale: usize, + json_shape: &str, + format: BenchFormat, +) -> Result<(usize, u64)> { + match format { + BenchFormat::Atif => write_atif_corpus(path, scale, json_shape), + BenchFormat::Actf => write_actf_corpus(path, scale, json_shape), + } } -fn write_corpus(path: &std::path::Path, scale: usize, json_shape: &str) -> Result<(usize, u64)> { +fn write_atif_corpus( + path: &std::path::Path, + scale: usize, + json_shape: &str, +) -> Result<(usize, u64)> { let fixture_root = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/atif"); let fixture_names = [ "dialogue_10.json", @@ -213,6 +437,64 @@ fn write_corpus(path: &std::path::Path, scale: usize, json_shape: &str) -> Resul output.flush()?; Ok((fixtures.len() * scale, steps)) } + +fn write_actf_corpus( + path: &std::path::Path, + scale: usize, + json_shape: &str, +) -> Result<(usize, u64)> { + let fixture_root = + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/import_roundtrip"); + let fixture_names = [ + "protein-assembly_trimmed.actf.json", + "make-doom-for-mips_trimmed.actf.json", + ]; + let fixtures = fixture_names + .iter() + .map(|name| std::fs::read_to_string(fixture_root.join(name))) + .map(|result| result.map_err(anyhow::Error::from)) + .map(|result| { + result.and_then(|text| { + serde_json::from_str::(&text).map_err(Into::into) + }) + }) + .collect::>>()?; + let mut output = BufWriter::new(File::create(path)?); + let mut steps = 0_u64; + let array = json_shape == "array"; + let mut first = true; + if array { + output.write_all(b"[\n")?; + } + for copy in 0..scale { + for (index, fixture) in fixtures.iter().enumerate() { + let mut document = fixture.clone(); + let task_id = format!("stream-{copy:06}-{index:02}"); + document["task_id"] = serde_json::Value::String(task_id); + steps += document["attempts"] + .as_object() + .and_then(|attempts| attempts.values().next()) + .and_then(|attempt| attempt.get("trajectory")) + .and_then(|trajectory| trajectory.get("steps")) + .and_then(|steps| steps.as_array()) + .map_or(0, Vec::len) as u64; + if array && !first { + output.write_all(b",\n")?; + } + serde_json::to_writer(&mut output, &document)?; + if !array { + output.write_all(b"\n")?; + } + first = false; + } + } + if array { + output.write_all(b"\n]\n")?; + } + output.flush()?; + Ok((fixtures.len() * scale, steps)) +} + fn reset_allocations() { ALLOCATION_COUNT.store(0, Ordering::SeqCst); ALLOCATED_BYTES.store(0, Ordering::SeqCst); diff --git a/crates/persisting-pchronicle/src/agenticmd/codec.rs b/crates/persisting-pchronicle/src/agenticmd/codec.rs index e4e96e5a..3b52806f 100644 --- a/crates/persisting-pchronicle/src/agenticmd/codec.rs +++ b/crates/persisting-pchronicle/src/agenticmd/codec.rs @@ -1,8 +1,7 @@ //! `agenticmd` — best-effort Markdown view for humans and debugging. //! -//! It is intentionally not a canonical storage format. New writers use -//! Storyline-like identity fields (`session_id`, `agent_id`, `source`, -//! `step_id`); readers retain aliases for older capture documents. +//! It is intentionally not a canonical storage format. Storyline metadata +//! carries semantics; block headers provide a readable, editable view. //! ```text //! --- //! format: persisting # logical name in pChronicle: agenticmd @@ -46,7 +45,7 @@ pub struct MarkdownBlock { } impl MarkdownBlock { - /// Legacy presentation role, derived from Storyline `source` when absent. + /// Human-facing presentation role derived from Storyline `source`. pub fn role(&self) -> Option<&str> { if let Some(role) = self.header.fields.get("role").and_then(|v| v.as_str()) { return Some(role); @@ -71,12 +70,6 @@ impl MarkdownBlock { _ => Some("system"), } } - - pub fn step_id(&self) -> Option { - ["step_id", "id", "seq"] - .iter() - .find_map(|key| self.header.fields.get(*key).and_then(|v| v.as_i64())) - } } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] @@ -85,10 +78,6 @@ pub struct MarkdownDocument { pub format: String, /// Frontmatter `format:` value (usually `persisting`). pub frontmatter_format: String, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub session_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub agent_id: Option, #[serde(default)] pub frontmatter: BTreeMap, pub blocks: Vec, @@ -99,8 +88,6 @@ impl MarkdownDocument { Self { format: AGENTICMD_FORMAT_NAME.into(), frontmatter_format: AGENTICMD_FRONTMATTER_FORMAT.into(), - session_id: None, - agent_id: None, frontmatter: BTreeMap::new(), blocks, } @@ -115,25 +102,6 @@ pub fn parse_agenticmd_document(input: &str) -> InputResult { if let Some(fmt) = doc.frontmatter.get("format").and_then(Value::as_str) { doc.frontmatter_format = fmt.to_string(); } - doc.session_id = doc - .frontmatter - .get("session_id") - .or_else(|| doc.frontmatter.get("session")) - .and_then(Value::as_str) - .map(str::to_string); - doc.agent_id = doc - .frontmatter - .get("agent_id") - .and_then(Value::as_str) - .or_else(|| doc.frontmatter.get("agent").and_then(Value::as_str)) - .or_else(|| { - doc.frontmatter - .get("agent") - .and_then(Value::as_object) - .and_then(|agent| agent.get("id")) - .and_then(Value::as_str) - }) - .map(str::to_string); Ok(doc) } diff --git a/crates/persisting-pchronicle/src/agenticmd/convert.rs b/crates/persisting-pchronicle/src/agenticmd/convert.rs index 655bcf89..720c42e6 100644 --- a/crates/persisting-pchronicle/src/agenticmd/convert.rs +++ b/crates/persisting-pchronicle/src/agenticmd/convert.rs @@ -1,19 +1,24 @@ //! agenticmd ⇄ storyline. //! -//! AgenticMD is a human/debugging view. Conversion prefers Storyline field -//! names while retaining legacy aliases for older capture documents. +//! AgenticMD is a human/debugging view backed by authoritative Storyline +//! metadata and readable message blocks. use std::collections::BTreeMap; use serde_json::{json, Value}; -use crate::formats::storyline::{StorylineAgent, StorylineDocument, StorylineTurn}; +use crate::formats::storyline::{StorylineDocument, StorylineTurn}; +use crate::formats::unknown_fields::{ + canonical_source_document_id, normalize_agenticmd_unknown_pointer, restore_json_pointer, + PointerWrite, UnknownFieldLimits, +}; use crate::{InputIssue, InputResult, Result}; use super::codec::{ encode_agenticmd_block, encode_agenticmd_preamble, parse_agenticmd_document, MarkdownBlock, MarkdownDocument, MarkdownHeader, AGENTICMD_FRONTMATTER_FORMAT, }; +use super::validate::{validate_agenticmd_storyline, validate_agenticmd_unknown_pointer}; const STORYLINE_METADATA_KEY: &str = "storyline"; const MESSAGE_ENCODING_KEY: &str = "message_encoding"; @@ -21,9 +26,10 @@ const MESSAGE_ENCODING_KEY: &str = "message_encoding"; /// Parse AgenticMD into its authoritative Storyline model. pub fn parse_agenticmd(input: &str) -> InputResult { let document = parse_agenticmd_document(input)?; - let Some(metadata) = document.frontmatter.get(STORYLINE_METADATA_KEY) else { - return agenticmd_to_storyline(&document); - }; + let metadata = document + .frontmatter + .get(STORYLINE_METADATA_KEY) + .ok_or_else(|| InputIssue::invalid("missing authoritative Storyline metadata"))?; let mut story = metadata .as_object() .cloned() @@ -64,36 +70,63 @@ pub fn parse_agenticmd(input: &str) -> InputResult { "turns".into(), serde_json::to_value(&turns).map_err(|error| InputIssue::invalid(error.to_string()))?, ); - let document = serde_json::from_value::(Value::Object(story)) + let mut story = serde_json::from_value::(Value::Object(story)) .map_err(|error| InputIssue::invalid(error.to_string()).at("frontmatter.storyline"))?; - document.validate()?; - Ok(document) + validate_agenticmd_storyline(&story)?; + capture_agenticmd_unknown_fields(&document, &mut story)?; + Ok(story) } /// Encode a Storyline as its human-readable AgenticMD representation. pub fn encode_agenticmd(story: &StorylineDocument) -> Result { - story.validate()?; - let mut output = encode_storyline_preamble(story)?; - for turn in &story.turns { - output.push_str(&encode_agenticmd_block(&storyline_turn_block(turn, None)?)?); + validate_agenticmd_storyline(story)?; + let frontmatter = storyline_frontmatter(story)?; + let blocks = story + .turns + .iter() + .map(|turn| storyline_turn_block(turn, None)) + .collect::>>()?; + let mut logical_document = json!({ + "frontmatter": frontmatter, + "blocks": blocks, + }); + restore_agenticmd_unknown_fields(story, &mut logical_document)?; + + let frontmatter = + serde_json::from_value::>(logical_document["frontmatter"].clone())?; + let blocks = serde_json::from_value::>(logical_document["blocks"].clone())?; + + let mut output = encode_agenticmd_preamble(&frontmatter)?; + for block in blocks { + output.push_str(&encode_agenticmd_block(&block)?); } Ok(output) } pub(super) fn encode_storyline_preamble(story: &StorylineDocument) -> Result { - let mut metadata = serde_json::to_value(story)? + validate_agenticmd_storyline(story)?; + encode_agenticmd_preamble(&storyline_frontmatter(story)?) +} + +fn storyline_frontmatter(story: &StorylineDocument) -> Result> { + // Native fields are written back to their Markdown locations below. The + // existing Storyline metadata remains the carrier for all foreign sources. + let mut metadata_story = story.clone(); + metadata_story.unknown_fields.sources.remove("agenticmd"); + metadata_story.unknown_key_counts.remove("agenticmd"); + + let mut metadata = serde_json::to_value(metadata_story)? .as_object() .cloned() .ok_or_else(|| anyhow::anyhow!("serialized Storyline must be an object"))?; metadata.remove("turns"); - let frontmatter: BTreeMap = BTreeMap::from([ + Ok(BTreeMap::from([ ( "format".into(), Value::String(AGENTICMD_FRONTMATTER_FORMAT.into()), ), (STORYLINE_METADATA_KEY.into(), Value::Object(metadata)), - ]); - encode_agenticmd_preamble(&frontmatter) + ])) } pub(super) fn storyline_turn_block( @@ -128,88 +161,104 @@ pub(super) fn storyline_turn_block( }) } -fn agenticmd_to_storyline(doc: &MarkdownDocument) -> InputResult { - let session_id = doc.session_id.clone().unwrap_or_else(|| "unknown".into()); - let agent_id = doc.agent_id.clone().unwrap_or_else(|| "unknown".into()); - - let mut turns = Vec::new(); - for (i, block) in doc.blocks.iter().enumerate() { - let id = block.step_id().unwrap_or((i as i64) + 1); - let source = block.source().unwrap_or("system"); - let model = block - .header - .fields - .get("model") - .and_then(|v| v.as_str()) - .map(str::to_string); - let latency_ms = block - .header - .fields - .get("latency_ms") - .and_then(|v| v.as_i64()); - let ttft_ms = block.header.fields.get("ttft_ms").and_then(|v| v.as_i64()); - let kind = block - .header - .fields - .get("kind") - .and_then(|v| v.as_str()) - .map(str::to_string); - let timestamp = block - .header - .fields - .get("timestamp") - .and_then(|v| v.as_str()) - .map(str::to_string); - - turns.push(StorylineTurn { - id, - kind, - timestamp, - source: source.into(), - message: Value::String(block.body.clone()), - reasoning_content: None, - reasoning_effort: None, - tool_calls: None, - observation: None, - metrics: None, - model_name: model, - llm_call_count: if source == "agent" { Some(1) } else { None }, - is_copied_context: None, - latency_ms, - ttft_ms, - extra: None, - }); +fn capture_agenticmd_unknown_fields( + document: &MarkdownDocument, + story: &mut StorylineDocument, +) -> InputResult<()> { + let source_document_id = agenticmd_source_document_id(document)?; + for (key, value) in &document.frontmatter { + if is_consumed_frontmatter_field(key) { + continue; + } + story.unknown_fields.insert( + "agenticmd", + &source_document_id, + format!("/frontmatter/{}", encode_pointer_token(key)), + value.clone(), + )?; } - Ok(StorylineDocument { - schema_version: None, - run_id: None, - trajectory_id: None, - attempt_id: None, - session_id, - agent: StorylineAgent { - id: agent_id.clone(), - name: Some(agent_id), - version: None, - model_name: None, - tool_definitions: None, - extra: None, - }, - parent: None, - child_session_ids: None, - notes: None, - final_metrics: None, - continued_trajectory_ref: None, - extra: None, - presence: Default::default(), - turns, - }) + for (index, block) in document.blocks.iter().enumerate() { + for (key, value) in &block.header.fields { + if is_consumed_header_field(key) { + continue; + } + story.unknown_fields.insert( + "agenticmd", + &source_document_id, + format!("/blocks/{index}/header/{}", encode_pointer_token(key)), + value.clone(), + )?; + } + } + + let recomputed_counts = story.unknown_fields.validate_with( + UnknownFieldLimits::default(), + normalize_agenticmd_unknown_pointer, + )?; + match recomputed_counts.get("agenticmd") { + Some(counts) => { + story + .unknown_key_counts + .insert("agenticmd".into(), counts.clone()); + } + None => { + story.unknown_key_counts.remove("agenticmd"); + } + } + Ok(()) +} + +fn agenticmd_source_document_id(document: &MarkdownDocument) -> InputResult { + let mut source = + serde_json::to_value(document).map_err(|error| InputIssue::invalid(error.to_string()))?; + source + .get_mut("frontmatter") + .and_then(Value::as_object_mut) + .expect("serialized AgenticMD document has object frontmatter") + .remove(STORYLINE_METADATA_KEY); + canonical_source_document_id(&source).map_err(|error| InputIssue::invalid(error.to_string())) +} + +fn restore_agenticmd_unknown_fields( + story: &StorylineDocument, + logical_document: &mut Value, +) -> Result<()> { + let Some(source) = story.unknown_fields.sources.get("agenticmd") else { + return Ok(()); + }; + for (pointer, value) in &source.fields { + validate_agenticmd_unknown_pointer(pointer)?; + restore_json_pointer( + logical_document, + pointer, + value.clone(), + PointerWrite::InsertOnly, + )?; + } + Ok(()) +} + +fn is_consumed_frontmatter_field(key: &str) -> bool { + matches!(key, "format" | STORYLINE_METADATA_KEY) +} + +fn is_consumed_header_field(key: &str) -> bool { + matches!( + key, + "source" | "step_id" | MESSAGE_ENCODING_KEY | STORYLINE_METADATA_KEY + ) +} + +fn encode_pointer_token(token: &str) -> String { + token.replace('~', "~0").replace('/', "~1") } #[cfg(test)] mod tests { use super::{encode_agenticmd, parse_agenticmd}; - use crate::{FieldPresence, StoryLink, StorylineDocument, StorylineToolCall, StorylineTurn}; + use crate::formats::unknown_fields::normalize_agenticmd_unknown_pointer; + use crate::{StoryLink, StorylineDocument, StorylineToolCall, StorylineTurn}; use serde_json::json; #[test] @@ -245,7 +294,7 @@ mod tests { tool_call_id: "call-1".into(), function_name: "lookup".into(), arguments: json!({"q":"x"}), - result: FieldPresence::Null, + result: None, duration_ms: Some(12), extra: Some(json!({"provider":"test"})), }]), @@ -265,4 +314,134 @@ mod tests { let restored = parse_agenticmd(&markdown).unwrap(); assert_eq!(restored, story); } + + #[test] + fn agenticmd_frontmatter_carries_unknown_sources() { + let mut story = StorylineDocument::new("s", "a"); + story + .unknown_fields + .insert("atif", "source", "/vendor", json!(7)) + .unwrap(); + story.refresh_unknown_key_counts().unwrap(); + + let encoded = encode_agenticmd(&story).unwrap(); + let decoded = parse_agenticmd(&encoded).unwrap(); + + assert_eq!(decoded.unknown_fields, story.unknown_fields); + assert_eq!(decoded.unknown_key_counts, story.unknown_key_counts); + } + + #[test] + fn agenticmd_captures_and_restores_native_unknown_fields_at_logical_pointers() { + let input = r#"--- +format: persisting +storyline: + session: s + agent: + id: a +vendor_top: 7 +--- + + + +hi +"#; + + let parsed = parse_agenticmd(input).unwrap(); + let fields = &parsed.unknown_fields.sources["agenticmd"].fields; + assert_eq!(fields["/frontmatter/vendor_top"], json!(7)); + assert_eq!( + fields["/blocks/0/header/vendor_header"], + serde_json::Value::Null + ); + assert_eq!( + parsed.unknown_key_counts["agenticmd"]["/blocks/*/header/vendor_header"], + 1 + ); + + let encoded = encode_agenticmd(&parsed).unwrap(); + let restored = parse_agenticmd(&encoded).unwrap(); + assert_eq!(restored, parsed); + } + + #[test] + fn agenticmd_rejects_native_unknown_field_collisions() { + let mut story = StorylineDocument::new("s", "a"); + story + .unknown_fields + .insert( + "agenticmd", + "source", + "/frontmatter/format", + json!("vendor"), + ) + .unwrap(); + story.refresh_unknown_key_counts().unwrap(); + + assert!(encode_agenticmd(&story).is_err()); + } + + #[test] + fn agenticmd_rejects_native_block_header_collisions() { + let mut story = StorylineDocument::new("s", "a"); + story.turns.push(StorylineTurn { + id: 1, + kind: None, + timestamp: None, + source: "user".into(), + message: json!("hello"), + reasoning_content: None, + reasoning_effort: None, + tool_calls: None, + observation: None, + metrics: None, + model_name: None, + llm_call_count: None, + is_copied_context: None, + latency_ms: None, + ttft_ms: None, + extra: None, + }); + story + .unknown_fields + .insert( + "agenticmd", + "source", + "/blocks/0/header/source", + json!("vendor"), + ) + .unwrap(); + story.unknown_key_counts = story + .unknown_fields + .validate_with( + crate::formats::unknown_fields::UnknownFieldLimits::default(), + normalize_agenticmd_unknown_pointer, + ) + .unwrap(); + + assert!(encode_agenticmd(&story).is_err()); + } + + #[test] + fn agenticmd_rejects_mismatched_serialized_unknown_key_counts() { + let input = r#"--- +format: persisting +storyline: + session: s + agent: + id: a + unknown_fields: + sources: + atif: + source_document_id: source + fields: + /vendor: 7 + unknown_key_counts: + atif: + /vendor: 2 +--- +"#; + + assert!(parse_agenticmd(input).is_err()); + } } diff --git a/crates/persisting-pchronicle/src/agenticmd/validate.rs b/crates/persisting-pchronicle/src/agenticmd/validate.rs index b253686e..330935fd 100644 --- a/crates/persisting-pchronicle/src/agenticmd/validate.rs +++ b/crates/persisting-pchronicle/src/agenticmd/validate.rs @@ -3,7 +3,11 @@ //! Field presence and semantic combinations are deliberately not validated: //! AgenticMD is a debugging view, not a protocol boundary. -use crate::{InputIssue, InputResult}; +use crate::formats::unknown_fields::{ + compute_unknown_key_counts, normalize_agenticmd_unknown_pointer, validate_json_pointer, + UnknownFieldLimits, +}; +use crate::{InputIssue, InputResult, StorylineDocument}; use super::codec::{MarkdownBlock, MarkdownHeader}; @@ -58,3 +62,47 @@ pub fn validate_agenticmd_block(block: &MarkdownBlock) -> InputResult<()> { validate_speaker(block_speaker(&block.header))?; Ok(()) } + +/// Native AgenticMD unknown fields use its logical `blocks` array for key counts, +/// unlike foreign sources whose pointers retain their own format semantics. +pub(crate) fn validate_agenticmd_storyline(story: &StorylineDocument) -> InputResult<()> { + let expected_counts = story.unknown_fields.validate_with( + UnknownFieldLimits::default(), + normalize_agenticmd_unknown_pointer, + )?; + if expected_counts != story.unknown_key_counts { + return Err(InputIssue::invalid( + "storyline unknown_key_counts do not match AgenticMD unknown_fields", + )); + } + + // Reuse Storyline's canonical identity and turn validation without making + // its format-neutral count normalizer authoritative for AgenticMD. + let mut common = story.clone(); + common.unknown_key_counts = compute_unknown_key_counts(&common.unknown_fields)?; + common.validate() +} + +/// Only these leaf locations can be restored into an AgenticMD document. +/// The pointer token itself may be empty or escaped; it names a frontmatter +/// or header field and is deliberately not interpreted here. +pub(crate) fn validate_agenticmd_unknown_pointer(pointer: &str) -> InputResult<()> { + validate_json_pointer(pointer)?; + let tokens = pointer.split('/').collect::>(); + let is_frontmatter_field = matches!(tokens.as_slice(), ["", "frontmatter", _]); + let is_header_field = + matches!(tokens.as_slice(), ["", "blocks", _, "header", _]) && is_array_index(tokens[2]); + if is_frontmatter_field || is_header_field { + Ok(()) + } else { + Err(InputIssue::invalid(format!( + "AgenticMD unknown-field pointer '{pointer}' must target /frontmatter/ or /blocks//header/" + ))) + } +} + +fn is_array_index(token: &str) -> bool { + !token.is_empty() + && !(token.len() > 1 && token.starts_with('0')) + && token.parse::().is_ok() +} diff --git a/crates/persisting-pchronicle/src/atif.rs b/crates/persisting-pchronicle/src/atif.rs index 0322b626..1617af90 100644 --- a/crates/persisting-pchronicle/src/atif.rs +++ b/crates/persisting-pchronicle/src/atif.rs @@ -5,70 +5,74 @@ //! `runs` / `steps` / `tool_calls` schema. use serde::{Deserialize, Serialize}; -use serde_json::Value; - -use crate::formats::storyline::FieldPresence; +use serde_json::{Map, Value}; /// Root ATIF trajectory document. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct AtifTrajectory { pub schema_version: String, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub session_id: FieldPresence, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub trajectory_id: FieldPresence, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub session_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub trajectory_id: Option, pub agent: AtifAgent, pub steps: Vec, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub notes: FieldPresence, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub final_metrics: FieldPresence, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub continued_trajectory_ref: FieldPresence, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub extra: FieldPresence, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub subagent_trajectories: FieldPresence>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub notes: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub final_metrics: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub continued_trajectory_ref: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub extra: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub subagent_trajectories: Option>, + #[serde(default, flatten, skip_serializing_if = "Map::is_empty")] + pub unknown: Map, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct AtifAgent { pub name: String, pub version: String, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub model_name: FieldPresence, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub tool_definitions: FieldPresence, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub extra: FieldPresence, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model_name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tool_definitions: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub extra: Option, + #[serde(default, flatten, skip_serializing_if = "Map::is_empty")] + pub unknown: Map, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct AtifStep { pub step_id: i64, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub timestamp: FieldPresence, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub timestamp: Option, pub source: String, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub model_name: FieldPresence, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub reasoning_effort: FieldPresence, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model_name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub reasoning_effort: Option, /// String or multimodal content-part array. pub message: Value, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub reasoning_content: FieldPresence, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub tool_calls: FieldPresence>, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub observation: FieldPresence, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub metrics: FieldPresence, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub extra: FieldPresence, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub llm_call_count: FieldPresence, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub is_copied_context: FieldPresence, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub reasoning_content: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tool_calls: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub observation: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub metrics: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub extra: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub llm_call_count: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub is_copied_context: Option, + #[serde(default, flatten, skip_serializing_if = "Map::is_empty")] + pub unknown: Map, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] @@ -76,11 +80,12 @@ pub struct AtifToolCall { pub tool_call_id: String, pub function_name: String, pub arguments: Value, - /// Inline result. ATIF distinguishes an omitted result from explicit null. - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub result: FieldPresence, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub extra: FieldPresence, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub result: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub extra: Option, + #[serde(default, flatten, skip_serializing_if = "Map::is_empty")] + pub unknown: Map, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] @@ -89,6 +94,7 @@ pub struct AtifObservation { } impl AtifTrajectory { + #[cfg(test)] pub fn from_json_str(s: &str) -> crate::InputResult { let traj: Self = serde_json::from_str(s) .map_err(|error| crate::InputIssue::invalid(error.to_string()))?; @@ -100,10 +106,10 @@ impl AtifTrajectory { /// /// Preference: `session_id` → `trajectory_id` → error. pub fn effective_session_id(&self) -> crate::InputResult<&str> { - if let Some(id) = self.session_id.value().filter(|s| !s.is_empty()) { + if let Some(id) = self.session_id.as_ref().filter(|s| !s.is_empty()) { return Ok(id); } - if let Some(id) = self.trajectory_id.value().filter(|s| !s.is_empty()) { + if let Some(id) = self.trajectory_id.as_ref().filter(|s| !s.is_empty()) { return Ok(id); } Err(crate::InputIssue::invalid( @@ -124,7 +130,7 @@ impl AtifTrajectory { if embedded { let trajectory_id = self .trajectory_id - .value() + .as_ref() .filter(|value| !value.is_empty()) .ok_or_else(|| { crate::InputIssue::invalid("embedded ATIF trajectory requires trajectory_id") @@ -136,8 +142,10 @@ impl AtifTrajectory { } } else { let _ = self.effective_session_id()?; - if let Some(trajectory_id) = - self.trajectory_id.value().filter(|value| !value.is_empty()) + if let Some(trajectory_id) = self + .trajectory_id + .as_ref() + .filter(|value| !value.is_empty()) { trajectory_ids.insert(trajectory_id.clone()); } @@ -163,7 +171,7 @@ impl AtifTrajectory { step.step_id ))); } - if let Some(calls) = step.tool_calls.value() { + if let Some(calls) = step.tool_calls.as_ref() { for call in calls { if call.tool_call_id.is_empty() { return Err(crate::InputIssue::invalid("tool_call_id must be non-empty")); @@ -177,7 +185,7 @@ impl AtifTrajectory { } } } - if let Some(children) = self.subagent_trajectories.value() { + if let Some(children) = self.subagent_trajectories.as_ref() { for child in children { child.validate_inner(true, trajectory_ids)?; } diff --git a/crates/persisting-pchronicle/src/convert/actf.rs b/crates/persisting-pchronicle/src/convert/actf.rs index b15f88a5..2283d00d 100644 --- a/crates/persisting-pchronicle/src/convert/actf.rs +++ b/crates/persisting-pchronicle/src/convert/actf.rs @@ -2,8 +2,10 @@ use std::collections::BTreeMap; +use anyhow::Context as _; use serde_json::{json, Map, Value}; +use crate::format::DocumentFormat; use crate::formats::actf::{ ActfAttempt, ActfDocument, ActfObservation, ActfStep, ActfToolCall, ActfTrajectory, ACTF_SCHEMA_VERSION, @@ -11,165 +13,44 @@ use crate::formats::actf::{ use crate::formats::storyline::{ StorylineAgent, StorylineDocument, StorylineToolCall, StorylineTurn, }; +use crate::formats::unknown_fields::{ + decode_json_pointer, normalize_actf_pointer, restore_json_pointer, + validate_unknown_fields_with, write_foreign_unknown_fields_envelope, CarrierBinding, + PointerWrite, UnknownFieldLimits, +}; use crate::Result; -const ACTF_EXTENSION_KEY: &str = "persisting.dev/actf/v1"; - -#[cfg(test)] -pub fn actf_to_storyline(document: &ActfDocument) -> Result { - let mut stories = actf_to_storylines(document)?; - if stories.len() != 1 { - anyhow::bail!( - "ACTF document contains {} attempts; use actf_to_storylines", - stories.len() - ); - } - Ok(stories.remove(0)) -} - -pub fn actf_to_storylines(document: &ActfDocument) -> Result> { +pub(crate) fn actf_to_storylines(document: &ActfDocument) -> Result> { document.validate()?; - let root_metadata = root_metadata(document)?; let multiple_attempts = document.attempts.len() > 1; - document + let mut stories = document .attempts .iter() .map(|(attempt_id, attempt)| { - attempt_to_storyline( - document, - attempt_id, - attempt, - &root_metadata, - multiple_attempts, - ) + attempt_to_storyline(document, attempt_id, attempt, multiple_attempts) }) - .collect() -} - -#[cfg(test)] -pub fn storyline_to_actf(story: &StorylineDocument) -> Result { - storylines_to_actf(std::slice::from_ref(story)) -} - -pub fn storylines_to_actf(stories: &[StorylineDocument]) -> Result { - if stories.is_empty() { - anyhow::bail!("ACTF conversion requires at least one Storyline"); - } - let residual_count = stories - .iter() - .filter(|story| residual(story).is_some()) - .count(); - if residual_count == 0 { - if stories.len() != 1 { - anyhow::bail!("synthesizing ACTF without residual metadata requires one Storyline"); - } - return synthesize_actf(&stories[0]); - } - if residual_count != stories.len() { - anyhow::bail!("cannot mix ACTF residual and unrelated Storylines"); - } - - let first = residual(&stories[0]) - .ok_or_else(|| anyhow::anyhow!("ACTF residual disappeared during conversion"))?; - let root_value = first - .get("root") - .and_then(Value::as_object) - .ok_or_else(|| anyhow::anyhow!("ACTF residual missing root metadata"))? - .clone(); - let mut attempts = Map::new(); - for story in stories { - story.validate()?; - let metadata = residual(story) - .ok_or_else(|| anyhow::anyhow!("ACTF residual disappeared during conversion"))?; - if metadata.get("root").and_then(Value::as_object) != Some(&root_value) { - anyhow::bail!("ACTF Storylines have conflicting root residual"); - } - let attempt_id = metadata - .get("attempt_id") - .and_then(Value::as_str) - .filter(|value| !value.is_empty()) - .ok_or_else(|| anyhow::anyhow!("ACTF residual missing attempt_id"))?; - let mut attempt = metadata - .get("attempt") - .and_then(Value::as_object) - .cloned() - .ok_or_else(|| anyhow::anyhow!("ACTF residual missing attempt metadata"))?; - let mut trajectory = metadata - .get("trajectory") - .and_then(Value::as_object) - .cloned() - .ok_or_else(|| anyhow::anyhow!("ACTF residual missing trajectory metadata"))?; - let steps = story - .turns - .iter() - .map(storyline_step_value) - .collect::>>()?; - trajectory.insert("steps".into(), Value::Array(steps)); - let metrics = story.final_metrics.as_ref().and_then(Value::as_object); - attempt.insert( - "correct".into(), - metrics - .and_then(|value| value.get("correct")) - .cloned() - .unwrap_or(Value::Bool(false)), - ); - attempt.insert( - "score".into(), - metrics - .and_then(|value| value.get("score")) - .cloned() - .unwrap_or(Value::Null), - ); - attempt.insert( - "status".into(), - metrics - .and_then(|value| value.get("status")) - .cloned() - .unwrap_or_else(|| Value::String("completed".into())), - ); - attempt.insert("trajectory".into(), Value::Object(trajectory)); - if attempts - .insert(attempt_id.to_string(), Value::Object(attempt)) - .is_some() - { - anyhow::bail!("duplicate ACTF attempt id '{attempt_id}'"); - } + .collect::>>()?; + for ((attempt_id, attempt), story) in document.attempts.iter().zip(stories.iter_mut()) { + capture_actf_unknowns(document, attempt_id, attempt, story)?; + story.unknown_key_counts = validate_unknown_fields_with( + &story.unknown_fields, + UnknownFieldLimits::default(), + normalize_actf_pointer, + )?; } + Ok(stories) +} - let mut root = root_value; - root.insert( - "task_id".into(), - Value::String( - stories[0] - .run_id - .clone() - .unwrap_or_else(|| stories[0].session_id.clone()), - ), - ); - root.insert( - "correct".into(), - stories[0] - .final_metrics - .as_ref() - .and_then(|value| value.get("task_correct")) - .cloned() - .unwrap_or(Value::Bool(false)), - ); - root.insert("attempts".into(), Value::Object(attempts)); - let document: ActfDocument = serde_json::from_value(Value::Object(root))?; - document.validate()?; - Ok(document) +pub(crate) fn storylines_to_actf(stories: &[StorylineDocument]) -> Result { + storylines_to_actf_pointer(stories) } fn attempt_to_storyline( document: &ActfDocument, attempt_id: &str, attempt: &ActfAttempt, - root_metadata: &Value, multiple_attempts: bool, ) -> Result { - let attempt_metadata = attempt_residual(attempt)?; - let trajectory_metadata = trajectory_residual(&attempt.trajectory)?; let mut turns = Vec::with_capacity(attempt.trajectory.steps.len()); for step in &attempt.trajectory.steps { let tool_calls = (!step.tools.is_empty()) @@ -187,9 +68,7 @@ fn attempt_to_storyline( } else { None }, - extra: Some(json!({ - ACTF_EXTENSION_KEY: tool_residual(call)?, - })), + extra: None, }) }) .collect::>>() @@ -232,9 +111,7 @@ fn attempt_to_storyline( is_copied_context: None, latency_ms: step.metric.llm_infer_ms.as_f64().map(|value| value as i64), ttft_ms: None, - extra: Some(json!({ - ACTF_EXTENSION_KEY: step_residual(step)?, - })), + extra: None, }); } @@ -247,7 +124,7 @@ fn attempt_to_storyline( schema_version: None, run_id: Some(document.task_id.clone()), trajectory_id: None, - attempt_id: None, + attempt_id: Some(attempt_id.to_string()), session_id, agent: StorylineAgent { id: "actf-agent".into(), @@ -267,21 +144,264 @@ fn attempt_to_storyline( "task_correct": document.correct, })), continued_trajectory_ref: None, - extra: Some(json!({ - ACTF_EXTENSION_KEY: { - "root": root_metadata, - "attempt_id": attempt_id, - "attempt": attempt_metadata, - "trajectory": trajectory_metadata, - } - })), - presence: Default::default(), + extra: None, + unknown_fields: Default::default(), + unknown_key_counts: Default::default(), turns, }) } +fn capture_actf_unknowns( + document: &ActfDocument, + attempt_id: &str, + attempt: &ActfAttempt, + story: &mut StorylineDocument, +) -> crate::InputResult<()> { + let source_id = &document.task_id; + let mut root = serde_json::to_value(document) + .map_err(|error| crate::InputIssue::invalid(error.to_string()))?; + let root = root + .as_object_mut() + .expect("ACTF root serializes as object"); + for key in ["task_id", "correct", "attempts"] { + root.remove(key); + } + insert_actf_map(story, source_id, "", root)?; + + let attempt_prefix = pointer_join("/attempts", attempt_id); + let mut attempt_value = serde_json::to_value(attempt) + .map_err(|error| crate::InputIssue::invalid(error.to_string()))?; + let attempt_map = attempt_value + .as_object_mut() + .expect("ACTF attempt is object"); + for key in ["correct", "score", "status", "trajectory"] { + attempt_map.remove(key); + } + insert_actf_map(story, source_id, &attempt_prefix, attempt_map)?; + + let trajectory_prefix = pointer_join(&attempt_prefix, "trajectory"); + let mut trajectory_value = serde_json::to_value(&attempt.trajectory) + .map_err(|error| crate::InputIssue::invalid(error.to_string()))?; + let trajectory_map = trajectory_value + .as_object_mut() + .expect("ACTF trajectory is object"); + trajectory_map.remove("steps"); + insert_actf_map(story, source_id, &trajectory_prefix, trajectory_map)?; + + for (step_index, step) in attempt.trajectory.steps.iter().enumerate() { + let step_prefix = pointer_join( + &pointer_join(&trajectory_prefix, "steps"), + &step_index.to_string(), + ); + let mut step_value = serde_json::to_value(step) + .map_err(|error| crate::InputIssue::invalid(error.to_string()))?; + let step_map = step_value.as_object_mut().expect("ACTF step is object"); + let assistant = step_map.remove("assistant_content"); + for key in ["step_id", "metric", "tools", "observation", "started_at"] { + step_map.remove(key); + } + insert_actf_map(story, source_id, &step_prefix, step_map)?; + if let Some(mut assistant) = assistant { + let assistant = assistant.as_object_mut().expect("ACTF assistant is object"); + for key in ["content", "reasoning_content", "tool_calls"] { + assistant.remove(key); + } + insert_actf_map( + story, + source_id, + &pointer_join(&step_prefix, "assistant_content"), + assistant, + )?; + } + for (call_index, call) in step.tools.iter().enumerate() { + capture_actf_tool( + story, + source_id, + &pointer_join( + &pointer_join(&step_prefix, "tools"), + &call_index.to_string(), + ), + call, + )?; + } + for (call_index, call) in step.assistant_content.tool_calls.iter().enumerate() { + capture_actf_tool( + story, + source_id, + &pointer_join( + &pointer_join( + &pointer_join(&step_prefix, "assistant_content"), + "tool_calls", + ), + &call_index.to_string(), + ), + call, + )?; + } + } + Ok(()) +} + +fn capture_actf_tool( + story: &mut StorylineDocument, + source_id: &str, + prefix: &str, + call: &ActfToolCall, +) -> crate::InputResult<()> { + if call.kind != "tool_use" { + story.unknown_fields.insert( + "actf", + source_id, + pointer_join(prefix, "type"), + Value::String(call.kind.clone()), + )?; + } + let mut unknown = call.extra.clone(); + for key in ["name", "input", "command"] { + unknown.remove(key); + } + insert_actf_map(story, source_id, prefix, &unknown) +} + +fn insert_actf_map( + story: &mut StorylineDocument, + source_id: &str, + prefix: &str, + fields: &Map, +) -> crate::InputResult<()> { + for (key, value) in fields { + story + .unknown_fields + .insert("actf", source_id, pointer_join(prefix, key), value.clone())?; + } + Ok(()) +} + +fn pointer_join(parent: &str, token: &str) -> String { + format!("{parent}/{}", token.replace('~', "~0").replace('/', "~1")) +} + +fn storylines_to_actf_pointer(stories: &[StorylineDocument]) -> Result { + if stories.is_empty() { + anyhow::bail!("ACTF conversion requires at least one Storyline"); + } + let task_id = stories[0] + .run_id + .clone() + .unwrap_or_else(|| stories[0].session_id.clone()); + let mut attempts = Map::new(); + let mut carriers = Vec::new(); + for (story_index, story) in stories.iter().enumerate() { + let attempt_id = story + .attempt_id + .as_deref() + .or_else(|| (stories.len() == 1).then_some("1")) + .ok_or_else(|| anyhow::anyhow!("ACTF multi-attempt Storyline requires attempt_id"))?; + let canonical = synthesize_actf(story)?; + let attempt = serde_json::to_value(&canonical.attempts["1"])?; + if attempts.insert(attempt_id.into(), attempt).is_some() { + anyhow::bail!("duplicate ACTF attempt id '{attempt_id}'"); + } + carriers.push(CarrierBinding { + story_index, + pointer: pointer_join("/attempts", attempt_id), + }); + } + let task_correct = stories[0] + .final_metrics + .as_ref() + .and_then(|metrics| metrics.get("task_correct")) + .cloned() + .unwrap_or(Value::Bool(false)); + let mut value = json!({ + "task_id": task_id, + "category": "unknown", + "k": stories.len(), + "correct": task_correct, + "attempts_tried": stories.len(), + "solved_at": Value::Null, + "attempts": attempts, + }); + + let mut source_id = None::; + let mut unknown_fields = BTreeMap::::new(); + let actf_sources = stories + .iter() + .filter_map(|story| story.unknown_fields.sources.get("actf")) + .collect::>(); + if !actf_sources.is_empty() && actf_sources.len() != stories.len() { + anyhow::bail!("cannot mix ACTF unknown fields and unrelated Storylines"); + } + for source in actf_sources { + if source_id + .as_ref() + .is_some_and(|id| id != &source.source_document_id) + { + anyhow::bail!("one ACTF document cannot merge multiple source documents"); + } + source_id = Some(source.source_document_id.clone()); + for (pointer, field_value) in &source.fields { + match unknown_fields.get(pointer) { + Some(existing) if existing != field_value => { + anyhow::bail!("ACTF unknown-field conflict at '{pointer}'") + } + Some(_) => {} + None => { + unknown_fields.insert(pointer.clone(), field_value.clone()); + } + } + } + } + for (pointer, field_value) in unknown_fields { + let write = if is_actf_source_owned(&pointer) { + PointerWrite::ReplaceSourceOwned + } else { + PointerWrite::InsertOnly + }; + restore_json_pointer(&mut value, &pointer, field_value, write) + .with_context(|| format!("restore ACTF unknown field '{pointer}'"))?; + } + write_foreign_unknown_fields_envelope(DocumentFormat::Actf, &mut value, stories, &carriers)?; + let document: ActfDocument = serde_json::from_value(value)?; + document.validate()?; + Ok(document) +} + +fn is_actf_source_owned(pointer: &str) -> bool { + let Ok(tokens) = decode_json_pointer(pointer) else { + return false; + }; + if tokens.len() == 1 { + return matches!( + tokens[0].as_str(), + "category" | "k" | "attempts_tried" | "solved_at" + ); + } + let Some(last) = tokens.last().map(String::as_str) else { + return false; + }; + matches!( + last, + "final_answer" + | "ground_truth" + | "error" + | "artifacts" + | "extra" + | "analysis_result" + | "meta" + | "schema_version" + | "started_at" + | "finished_at" + | "system_prompt" + | "user_content" + | "type" + ) +} + fn synthesize_actf(story: &StorylineDocument) -> Result { - story.validate()?; + if story.session_id.is_empty() || story.agent.id.is_empty() { + anyhow::bail!("invalid Storyline identity for ACTF conversion"); + } let epoch = "1970-01-01 00:00:00+00:00".to_string(); let started_at = story .turns @@ -376,48 +496,12 @@ fn storyline_step_value(turn: &StorylineTurn) -> Result { .unwrap_or_default() .iter() .map(|call| { - let metadata = call - .extra - .as_ref() - .and_then(|extra| extra.get(ACTF_EXTENSION_KEY)) - .and_then(Value::as_object); - let mut tool = metadata - .and_then(|value| value.get("residual")) - .and_then(Value::as_object) - .cloned() - .unwrap_or_default(); - let kind = metadata - .and_then(|value| value.get("kind")) - .and_then(Value::as_str) - .unwrap_or("tool_use"); - tool.insert("type".into(), Value::String(kind.into())); - tool.insert("id".into(), Value::String(call.tool_call_id.clone())); - if metadata - .and_then(|value| value.get("name_present")) - .and_then(Value::as_bool) - .unwrap_or(true) - { - tool.insert("name".into(), Value::String(call.function_name.clone())); - } - match metadata - .and_then(|value| value.get("arguments_key")) - .and_then(Value::as_str) - .unwrap_or("input") - { - "command" => { - let command = call - .arguments - .get("command") - .cloned() - .unwrap_or_else(|| call.arguments.clone()); - tool.insert("command".into(), command); - } - "none" => {} - _ => { - tool.insert("input".into(), call.arguments.clone()); - } - } - Value::Object(tool) + json!({ + "type": "tool_use", + "id": call.tool_call_id, + "name": call.function_name, + "input": call.arguments, + }) }) .collect::>(); let observations = turn @@ -429,33 +513,42 @@ fn storyline_step_value(turn: &StorylineTurn) -> Result { .flatten() .map(|result| { let mut extra = result.as_object().cloned().unwrap_or_default(); - extra.remove("source_call_id"); + let source_call_id = extra.remove("source_call_id"); + extra + .entry("type") + .or_insert_with(|| Value::String("tool_result".into())); + if let Some(source_call_id) = source_call_id { + extra.entry("tool_use_id").or_insert(source_call_id); + } Value::Object(extra) }) .collect::>(); - let metric = turn.metrics.clone().unwrap_or_else(|| { - json!({ - "prompt_tokens_len": 0, - "completion_tokens_len": 0, - "llm_infer_ms": turn.latency_ms.map_or(Value::Null, |value| json!(value)), - "env_action_ms": Value::Null, - "stop_reason": Value::Null, - }) - }); - let timestamp = turn - .timestamp - .clone() - .unwrap_or_else(|| "1970-01-01 00:00:00+00:00".into()); - let metadata = turn - .extra + let mut metric = turn + .metrics .as_ref() - .and_then(|extra| extra.get(ACTF_EXTENSION_KEY)) - .and_then(Value::as_object); - let mut assistant = metadata - .and_then(|value| value.get("assistant_content")) .and_then(Value::as_object) .cloned() .unwrap_or_default(); + let prompt_tokens = metric.get("prompt_tokens").cloned().unwrap_or(json!(0)); + let completion_tokens = metric.get("completion_tokens").cloned().unwrap_or(json!(0)); + let llm_infer_ms = metric + .get("total_latency_ms") + .cloned() + .or_else(|| turn.latency_ms.map(|value| json!(value))) + .unwrap_or(Value::Null); + let stop_reason = metric.get("finish_reason").cloned().unwrap_or(Value::Null); + metric.entry("prompt_tokens_len").or_insert(prompt_tokens); + metric + .entry("completion_tokens_len") + .or_insert(completion_tokens); + metric.entry("llm_infer_ms").or_insert(llm_infer_ms); + metric.entry("env_action_ms").or_insert(Value::Null); + metric.entry("stop_reason").or_insert(stop_reason); + let timestamp = turn + .timestamp + .clone() + .unwrap_or_else(|| "1970-01-01 00:00:00+00:00".into()); + let mut assistant = Map::new(); assistant.insert( "content".into(), Value::String(turn.message.as_str().unwrap_or("").to_string()), @@ -469,25 +562,17 @@ fn storyline_step_value(turn: &StorylineTurn) -> Result { let mut step = Map::new(); step.insert("step_id".into(), json!(turn.id)); step.insert("assistant_content".into(), Value::Object(assistant)); - step.insert("metric".into(), metric); + step.insert("metric".into(), Value::Object(metric)); step.insert("tools".into(), Value::Array(tools)); step.insert("observation".into(), Value::Array(observations)); - let timestamp_style = metadata - .and_then(|value| value.get("started_at_style")) - .and_then(Value::as_str); - let started_at = metadata - .and_then(|value| value.get("started_at_original")) - .and_then(Value::as_str) - .filter(|original| *original == timestamp) - .map(str::to_string) - .map_or_else(|| format_actf_timestamp(×tamp, timestamp_style), Ok)?; + let started_at = format_actf_timestamp(×tamp)?; step.insert("started_at".into(), Value::String(started_at)); - if let Some(residual) = metadata - .and_then(|value| value.get("step")) - .and_then(Value::as_object) - { - merge_residual(&mut step, residual, "step"); - } + step.entry("system_prompt") + .or_insert_with(|| Value::String(String::new())); + step.entry("user_content") + .or_insert_with(|| Value::String(String::new())); + step.entry("finished_at") + .or_insert_with(|| Value::String(timestamp.clone())); Ok(Value::Object(step)) } @@ -518,124 +603,9 @@ fn actf_observation_call_id(observation: &ActfObservation) -> Option<&str> { .and_then(Value::as_str) } -fn root_metadata(document: &ActfDocument) -> Result { - let mut value = serde_json::to_value(document)?; - let object = value - .as_object_mut() - .ok_or_else(|| anyhow::anyhow!("serialized ACTF document must be an object"))?; - for key in ["task_id", "correct", "attempts"] { - object.remove(key); - } - Ok(value) -} - -fn attempt_residual(attempt: &ActfAttempt) -> Result { - let mut value = serde_json::to_value(attempt)?; - let object = value - .as_object_mut() - .ok_or_else(|| anyhow::anyhow!("serialized ACTF attempt must be an object"))?; - for key in ["correct", "score", "status", "trajectory"] { - object.remove(key); - } - Ok(value) -} - -fn trajectory_residual(trajectory: &ActfTrajectory) -> Result { - let mut value = serde_json::to_value(trajectory)?; - value - .as_object_mut() - .ok_or_else(|| anyhow::anyhow!("serialized ACTF trajectory must be an object"))? - .remove("steps"); - Ok(value) -} - -fn step_residual(step: &ActfStep) -> Result { - let mut value = serde_json::to_value(step)?; - let object = value - .as_object_mut() - .ok_or_else(|| anyhow::anyhow!("serialized ACTF step must be an object"))?; - let mut assistant = object - .remove("assistant_content") - .and_then(|value| value.as_object().cloned()) - .ok_or_else(|| anyhow::anyhow!("serialized ACTF assistant_content must be an object"))?; - for key in ["content", "reasoning_content", "tool_calls"] { - assistant.remove(key); - } - for key in ["step_id", "metric", "tools", "observation", "started_at"] { - object.remove(key); - } - let mut residual = Map::new(); - residual.insert("step".into(), Value::Object(object.clone())); - residual.insert("assistant_content".into(), Value::Object(assistant)); - residual.insert( - "started_at_style".into(), - Value::String(timestamp_style(&step.started_at).into()), - ); - residual.insert( - "started_at_original".into(), - Value::String(step.started_at.clone()), - ); - Ok(Value::Object(residual)) -} - -fn tool_residual(call: &ActfToolCall) -> Result { - let name_present = call.extra.contains_key("name"); - let arguments_key = if call.extra.contains_key("input") { - "input" - } else if call.extra.contains_key("command") { - "command" - } else { - "none" - }; - let mut residual = call.extra.clone(); - residual.remove("name"); - residual.remove("input"); - residual.remove("command"); - Ok(json!({ - "kind": call.kind, - "name_present": name_present, - "arguments_key": arguments_key, - "residual": residual, - })) -} - -fn merge_residual(target: &mut Map, residual: &Map, scope: &str) { - for (key, value) in residual { - if target.contains_key(key) { - tracing::warn!( - source_format = "actf", - source_key = %key, - target_key = %key, - scope, - "ACTF residual conflicts with an authoritative Storyline field" - ); - continue; - } - target.insert(key.clone(), value.clone()); - } -} - -fn timestamp_style(value: &str) -> &'static str { - if value.contains(' ') { - "space-offset" - } else if value.ends_with('Z') { - "rfc3339-z" - } else { - "rfc3339-offset" - } -} - -fn format_actf_timestamp(value: &str, style: Option<&str>) -> Result { +fn format_actf_timestamp(value: &str) -> Result { let timestamp = chrono::DateTime::parse_from_rfc3339(value)?; - Ok(match style { - Some("space-offset") => timestamp.format("%Y-%m-%d %H:%M:%S%.f%:z").to_string(), - Some("rfc3339-offset") => timestamp.to_rfc3339(), - _ => timestamp.to_rfc3339_opts(chrono::SecondsFormat::AutoSi, true), - }) -} - -fn residual(story: &StorylineDocument) -> Option<&Map> { - story.extra.as_ref()?.get(ACTF_EXTENSION_KEY)?.as_object() + Ok(timestamp.format("%Y-%m-%d %H:%M:%S%.f%:z").to_string()) } #[cfg(test)] @@ -645,6 +615,20 @@ mod tests { #[cfg(feature = "lance-store")] use crate::store::StorylineLanceStore; + fn actf_to_storyline(document: &ActfDocument) -> Result { + let mut stories = actf_to_storylines(document)?; + anyhow::ensure!( + stories.len() == 1, + "test fixture contains {} ACTF attempts", + stories.len() + ); + Ok(stories.remove(0)) + } + + fn storyline_to_actf(story: &StorylineDocument) -> Result { + storylines_to_actf(std::slice::from_ref(story)) + } + const FIXTURE: &str = r#"{ "task_id":"task-1","category":"software-engineering","k":1, "correct":false,"attempts_tried":1,"solved_at":null, @@ -675,7 +659,7 @@ mod tests { } #[test] - fn actf_residual_preserves_unknowns_but_storyline_fields_are_authoritative() { + fn actf_unknown_fields_preserve_values_but_storyline_fields_are_authoritative() { let mut value: Value = serde_json::from_str(FIXTURE).unwrap(); value["root_unknown"] = Value::Null; value["attempts"]["1"]["attempt_unknown"] = json!([3, 2, 1]); @@ -717,6 +701,26 @@ mod tests { assert_eq!(step["tools"][0]["tool_unknown"], Value::Null); } + #[test] + fn actf_unknown_fields_use_namespaced_exact_paths() { + let mut value: Value = serde_json::from_str(FIXTURE).unwrap(); + value["root_unknown"] = Value::Null; + value["attempts"]["1"]["trajectory"]["steps"][0]["step_unknown"] = json!({"x": 1}); + value["attempts"]["1"]["trajectory"]["steps"][0]["0"] = json!("literal"); + let document: ActfDocument = serde_json::from_value(value).unwrap(); + let stories = actf_to_storylines(&document).unwrap(); + let source = &stories[0].unknown_fields.sources["actf"]; + assert_eq!(source.fields["/root_unknown"], Value::Null); + assert_eq!( + source.fields["/attempts/1/trajectory/steps/0/step_unknown"], + json!({"x": 1}) + ); + assert_eq!( + stories[0].unknown_key_counts["actf"]["/attempts/1/trajectory/steps/*/0"], + 1 + ); + } + #[test] fn multiple_attempts_roundtrip_as_multiple_storylines() { let mut document = parse_actf_document(FIXTURE).unwrap(); @@ -733,6 +737,46 @@ mod tests { assert_eq!(storylines_to_actf(&stories).unwrap(), document); } + #[test] + fn synthesis_completes_partial_metrics_and_normalizes_observations() { + let mut story = StorylineDocument::new("session", "agent"); + story.turns.push(StorylineTurn { + id: 1, + kind: Some("autonomous".into()), + timestamp: None, + source: "agent".into(), + message: json!("done"), + reasoning_content: None, + reasoning_effort: None, + tool_calls: Some(vec![StorylineToolCall { + tool_call_id: "call-1".into(), + function_name: "inspect".into(), + arguments: json!({"path": "/tmp"}), + result: None, + duration_ms: None, + extra: None, + }]), + observation: Some(json!({ + "results": [{"source_call_id": "call-1", "content": "ok"}] + })), + metrics: Some(json!({"reward": 1.0})), + model_name: None, + llm_call_count: Some(1), + is_copied_context: None, + latency_ms: None, + ttft_ms: None, + extra: None, + }); + + let document = storyline_to_actf(&story).unwrap(); + let step = &document.attempts["1"].trajectory.steps[0]; + assert_eq!(step.observation[0].kind, "tool_result"); + assert_eq!(step.observation[0].extra["tool_use_id"], "call-1"); + assert_eq!(step.metric.prompt_tokens_len, 0); + assert_eq!(step.metric.completion_tokens_len, 0); + assert_eq!(step.metric.extra["reward"], 1.0); + } + #[cfg(feature = "lance-store")] #[tokio::test] async fn actf_lance_import_and_restore_is_lossless() { diff --git a/crates/persisting-pchronicle/src/convert/atif.rs b/crates/persisting-pchronicle/src/convert/atif.rs index 188b535d..0b4a0cb9 100644 --- a/crates/persisting-pchronicle/src/convert/atif.rs +++ b/crates/persisting-pchronicle/src/convert/atif.rs @@ -1,17 +1,23 @@ //! ATIF ⇄ storyline. use crate::atif::{AtifAgent, AtifObservation, AtifStep, AtifToolCall, AtifTrajectory}; +use crate::format::DocumentFormat; use crate::formats::storyline::{ - FieldPresence, PresenceState, StoryLink, StorylineAgent, StorylineAgentField, - StorylineCollectionShape, StorylineDocument, StorylinePresence, StorylineRootField, - StorylineToolCall, StorylineTurn, StorylineTurnField, + StoryLink, StorylineAgent, StorylineDocument, StorylineToolCall, StorylineTurn, +}; +use crate::formats::unknown_fields::{ + attach_carried_unknown_fields, canonical_source_document_id, restore_json_pointer, + take_unknown_fields_envelope, validate_unknown_fields, write_foreign_unknown_fields_envelope, + CarrierBinding, PointerWrite, UnknownFieldLimits, }; use anyhow::Context as _; +use serde_json::{Map, Value}; +use std::collections::{BTreeMap, HashMap, HashSet}; use crate::Result; -fn timing_from_metrics(metrics: &FieldPresence) -> (Option, Option) { - let Some(m) = metrics.value() else { +fn timing_from_metrics(metrics: &Option) -> (Option, Option) { + let Some(m) = metrics.as_ref() else { return (None, None); }; let latency = m @@ -25,57 +31,11 @@ fn timing_from_metrics(metrics: &FieldPresence) -> (Option bool; -} - -impl NullField for FieldPresence { - fn is_null_field(&self) -> bool { - self.is_null() - } -} - -fn presence_state(field: &FieldPresence) -> PresenceState { - match field { - FieldPresence::Missing => PresenceState::Missing, - FieldPresence::Null => PresenceState::Null, - FieldPresence::Value(_) => PresenceState::Value, - } -} - -fn mark_root_null( - presence: &mut StorylinePresence, - field: StorylineRootField, - value: &FieldPresence, -) { - if value.is_null() { - presence.root_nulls.insert(field); - } -} - -fn mark_agent_null( - presence: &mut StorylinePresence, - field: StorylineAgentField, - value: &FieldPresence, -) { - if value.is_null() { - presence.agent_nulls.insert(field); - } -} - -fn field_from_option(value: Option, explicit_null: bool) -> FieldPresence { - match value { - Some(value) => FieldPresence::Value(value), - None if explicit_null => FieldPresence::Null, - None => FieldPresence::Missing, - } -} - #[cfg(test)] pub fn atif_to_storyline(traj: &AtifTrajectory) -> Result { if traj .subagent_trajectories - .value() + .as_ref() .is_some_and(|children| !children.is_empty()) { anyhow::bail!("embedded ATIF subagent trajectories require atif_to_storylines"); @@ -87,29 +47,54 @@ pub fn atif_to_storyline(traj: &AtifTrajectory) -> Result { /// /// Embedded subagents retain input order through each parent's `children` /// list. Missing child `session_id` values inherit the parent's effective -/// storage identity while their original presence remains explicit. +/// storage identity. +#[cfg(test)] pub fn atif_to_storylines(traj: &AtifTrajectory) -> Result> { + let source_id = canonical_source_document_id(&serde_json::to_value(traj)?)?; + atif_to_storylines_with_source(traj, &source_id).map(|(stories, _)| stories) +} + +fn atif_to_storylines_with_source( + traj: &AtifTrajectory, + source_document_id: &str, +) -> Result<(Vec, Vec)> { fn visit( trajectory: &AtifTrajectory, parent_key: Option<&str>, inherited_session_id: Option<&str>, + source_document_id: &str, + source_pointer: &str, output: &mut Vec, + carriers: &mut Vec, ) -> Result<()> { - let story = atif_to_storyline_node(trajectory, parent_key, inherited_session_id)?; + let mut story = atif_to_storyline_node(trajectory, parent_key, inherited_session_id)?; + capture_atif_unknowns(trajectory, source_document_id, source_pointer, &mut story)?; let parent_key = story .trajectory_id .as_deref() .unwrap_or(story.session_id.as_str()) .to_string(); let inherited_session_id = story.session_id.clone(); + let story_index = output.len(); output.push(story); - if let Some(children) = trajectory.subagent_trajectories.value() { - for child in children { + carriers.push(CarrierBinding { + story_index, + pointer: source_pointer.to_string(), + }); + if let Some(children) = trajectory.subagent_trajectories.as_ref() { + for (index, child) in children.iter().enumerate() { + let child_pointer = pointer_join( + &pointer_join(source_pointer, "subagent_trajectories"), + &index.to_string(), + ); visit( child, Some(&parent_key), Some(&inherited_session_id), + source_document_id, + &child_pointer, output, + carriers, )?; } } @@ -117,28 +102,47 @@ pub fn atif_to_storylines(traj: &AtifTrajectory) -> Result Result> { - if ordinal < 0 { - anyhow::bail!("ATIF collection ordinal cannot be negative"); - } - let mut stories = atif_to_storylines(traj)?; +pub(crate) fn atif_value_to_storylines( + mut value: Value, +) -> crate::InputResult> { + let envelope = take_unknown_fields_envelope(&mut value)?; + let source_document_id = canonical_source_document_id(&value) + .map_err(|error| crate::InputIssue::invalid(error.to_string()))?; + let trajectory: AtifTrajectory = serde_json::from_value(value) + .map_err(|error| crate::InputIssue::invalid(error.to_string()))?; + trajectory.validate()?; + let (mut stories, carriers) = atif_to_storylines_with_source(&trajectory, &source_document_id) + .map_err(|error| crate::InputIssue::invalid(error.to_string()))?; + attach_carried_unknown_fields( + envelope, + &carriers, + &mut stories, + UnknownFieldLimits::default(), + )?; for story in &mut stories { - story.presence.collection_shape = Some(shape); - story.presence.collection_ordinal = Some(ordinal); + story.unknown_key_counts = + validate_unknown_fields(&story.unknown_fields, UnknownFieldLimits::default())?; } Ok(stories) } +pub(crate) fn atif_collection_to_storylines(value: Value) -> Result> { + atif_value_to_storylines(value).map_err(anyhow::Error::from) +} + fn atif_to_storyline_node( traj: &AtifTrajectory, parent_key: Option<&str>, @@ -146,25 +150,23 @@ fn atif_to_storyline_node( ) -> Result { let session_id = traj .session_id - .value() - .map(String::as_str) + .as_deref() .filter(|value| !value.is_empty()) .or(inherited_session_id) .or_else(|| { traj.trajectory_id - .value() - .map(String::as_str) + .as_deref() .filter(|value| !value.is_empty()) }) .ok_or_else(|| anyhow::anyhow!("ATIF trajectory requires an effective storage identity"))? .to_string(); - let child_ids = traj.subagent_trajectories.value().map(|children| { + let child_ids = traj.subagent_trajectories.as_ref().map(|children| { children .iter() .map(|child| { child .trajectory_id - .value() + .as_ref() .filter(|value| !value.is_empty()) .ok_or_else(|| anyhow::anyhow!("embedded ATIF subagent requires trajectory_id")) .cloned() @@ -173,94 +175,24 @@ fn atif_to_storyline_node( }); let child_ids = child_ids.transpose()?; - let mut presence = StorylinePresence { - session_id: presence_state(&traj.session_id), - ..StorylinePresence::default() - }; - mark_root_null( - &mut presence, - StorylineRootField::TrajectoryId, - &traj.trajectory_id, - ); - mark_root_null(&mut presence, StorylineRootField::Notes, &traj.notes); - mark_root_null( - &mut presence, - StorylineRootField::FinalMetrics, - &traj.final_metrics, - ); - mark_root_null( - &mut presence, - StorylineRootField::ContinuedTrajectoryRef, - &traj.continued_trajectory_ref, - ); - mark_root_null(&mut presence, StorylineRootField::Extra, &traj.extra); - mark_root_null( - &mut presence, - StorylineRootField::SubagentTrajectories, - &traj.subagent_trajectories, - ); - mark_agent_null( - &mut presence, - StorylineAgentField::ModelName, - &traj.agent.model_name, - ); - mark_agent_null( - &mut presence, - StorylineAgentField::ToolDefinitions, - &traj.agent.tool_definitions, - ); - mark_agent_null(&mut presence, StorylineAgentField::Extra, &traj.agent.extra); - let mut turns = Vec::new(); for step in &traj.steps { - for (field, value) in [ - ( - StorylineTurnField::Timestamp, - &step.timestamp as &dyn NullField, - ), - (StorylineTurnField::ModelName, &step.model_name), - (StorylineTurnField::ReasoningEffort, &step.reasoning_effort), - ( - StorylineTurnField::ReasoningContent, - &step.reasoning_content, - ), - (StorylineTurnField::ToolCalls, &step.tool_calls), - (StorylineTurnField::Observation, &step.observation), - (StorylineTurnField::Metrics, &step.metrics), - (StorylineTurnField::Extra, &step.extra), - (StorylineTurnField::LlmCallCount, &step.llm_call_count), - (StorylineTurnField::IsCopiedContext, &step.is_copied_context), - ] { - if value.is_null_field() { - presence - .turn_nulls - .entry(step.step_id) - .or_default() - .insert(field); - } - } - - let tool_calls = step.tool_calls.value().map(|calls| { + let tool_calls = step.tool_calls.as_ref().map(|calls| { calls .iter() .map(|c| { let duration_ms = c .extra - .value() + .as_ref() .and_then(|x| x.get("duration_ms")) .and_then(|v| v.as_i64()); - if c.extra.is_null() { - presence - .tool_call_extra_nulls - .insert(c.tool_call_id.clone()); - } StorylineToolCall { tool_call_id: c.tool_call_id.clone(), function_name: c.function_name.clone(), arguments: c.arguments.clone(), result: c.result.clone(), duration_ms, - extra: c.extra.clone().into_option(), + extra: c.extra.clone(), } }) .collect::>() @@ -271,24 +203,24 @@ fn atif_to_storyline_node( let mut turn = StorylineTurn { id: step.step_id, kind: None, - timestamp: step.timestamp.clone().into_option(), + timestamp: step.timestamp.clone(), source: step.source.clone(), message: step.message.clone(), - reasoning_content: step.reasoning_content.clone().into_option(), - reasoning_effort: step.reasoning_effort.clone().into_option(), + reasoning_content: step.reasoning_content.clone(), + reasoning_effort: step.reasoning_effort.clone(), tool_calls, observation: step .observation - .value() + .as_ref() .map(serde_json::to_value) .transpose()?, - metrics: step.metrics.clone().into_option(), - model_name: step.model_name.clone().into_option(), - llm_call_count: step.llm_call_count.clone().into_option(), - is_copied_context: step.is_copied_context.clone().into_option(), + metrics: step.metrics.clone(), + model_name: step.model_name.clone(), + llm_call_count: step.llm_call_count, + is_copied_context: step.is_copied_context, latency_ms, ttft_ms, - extra: step.extra.clone().into_option(), + extra: step.extra.clone(), }; let derived = turn.effective_kind().to_string(); if !matches!( @@ -303,16 +235,16 @@ fn atif_to_storyline_node( Ok(StorylineDocument { schema_version: Some(traj.schema_version.clone()), run_id: None, - trajectory_id: traj.trajectory_id.clone().into_option(), + trajectory_id: traj.trajectory_id.clone(), attempt_id: None, session_id, agent: StorylineAgent { id: traj.agent.name.clone(), name: Some(traj.agent.name.clone()), version: Some(traj.agent.version.clone()), - model_name: traj.agent.model_name.clone().into_option(), - tool_definitions: traj.agent.tool_definitions.clone().into_option(), - extra: traj.agent.extra.clone().into_option(), + model_name: traj.agent.model_name.clone(), + tool_definitions: traj.agent.tool_definitions.clone(), + extra: traj.agent.extra.clone(), }, parent: parent_key.map(|parent_session_id| StoryLink { parent_session_id: parent_session_id.to_string(), @@ -321,15 +253,74 @@ fn atif_to_storyline_node( relation: "spawn".into(), }), child_session_ids: child_ids, - notes: traj.notes.clone().into_option(), - final_metrics: traj.final_metrics.clone().into_option(), - continued_trajectory_ref: traj.continued_trajectory_ref.clone().into_option(), - extra: traj.extra.clone().into_option(), - presence, + notes: traj.notes.clone(), + final_metrics: traj.final_metrics.clone(), + continued_trajectory_ref: traj.continued_trajectory_ref.clone(), + extra: traj.extra.clone(), + unknown_fields: Default::default(), + unknown_key_counts: Default::default(), turns, }) } +fn capture_atif_unknowns( + trajectory: &AtifTrajectory, + source_document_id: &str, + source_pointer: &str, + story: &mut StorylineDocument, +) -> crate::InputResult<()> { + insert_unknown_map( + story, + source_document_id, + source_pointer, + &trajectory.unknown, + )?; + insert_unknown_map( + story, + source_document_id, + &pointer_join(source_pointer, "agent"), + &trajectory.agent.unknown, + )?; + for (step_index, step) in trajectory.steps.iter().enumerate() { + let step_pointer = pointer_join( + &pointer_join(source_pointer, "steps"), + &step_index.to_string(), + ); + insert_unknown_map(story, source_document_id, &step_pointer, &step.unknown)?; + if let Some(calls) = step.tool_calls.as_ref() { + for (call_index, call) in calls.iter().enumerate() { + let call_pointer = pointer_join( + &pointer_join(&step_pointer, "tool_calls"), + &call_index.to_string(), + ); + insert_unknown_map(story, source_document_id, &call_pointer, &call.unknown)?; + } + } + } + Ok(()) +} + +fn insert_unknown_map( + story: &mut StorylineDocument, + source_document_id: &str, + parent: &str, + fields: &Map, +) -> crate::InputResult<()> { + for (key, value) in fields { + story.unknown_fields.insert( + "atif", + source_document_id, + pointer_join(parent, key), + value.clone(), + )?; + } + Ok(()) +} + +fn pointer_join(parent: &str, token: &str) -> String { + format!("{parent}/{}", token.replace('~', "~0").replace('/', "~1")) +} + #[cfg(test)] pub fn storyline_to_atif(story: &StorylineDocument) -> Result { if story @@ -346,7 +337,9 @@ fn storyline_to_atif_node( story: &StorylineDocument, embedded_children: Option>, ) -> Result { - story.validate()?; + if story.session_id.is_empty() || story.agent.id.is_empty() { + anyhow::bail!("invalid Storyline identity for ATIF conversion"); + } let mut steps = Vec::new(); for (step_index, turn) in story.turns.iter().enumerate() { let observation = turn @@ -369,17 +362,10 @@ fn storyline_to_atif_node( obj.insert("duration_ms".into(), serde_json::json!(ms)); } } - let extra = if story - .presence - .tool_call_extra_nulls - .contains(&c.tool_call_id) - && extra.as_object().is_some_and(|object| object.is_empty()) - { - FieldPresence::Null - } else if extra.as_object().map(|o| o.is_empty()).unwrap_or(true) { - FieldPresence::Missing + let extra = if extra.as_object().map(|o| o.is_empty()).unwrap_or(true) { + None } else { - FieldPresence::Value(extra) + Some(extra) }; AtifToolCall { tool_call_id: c.tool_call_id.clone(), @@ -387,6 +373,7 @@ fn storyline_to_atif_node( arguments: c.arguments.clone(), result: c.result.clone(), extra, + unknown: Map::new(), } }) .collect() @@ -403,72 +390,34 @@ fn storyline_to_atif_node( .or_insert(serde_json::json!(ms)); } } - let metrics = if turn.metrics.is_none() - && story - .presence - .turn_nulls - .get(&turn.id) - .is_some_and(|fields| fields.contains(&StorylineTurnField::Metrics)) - { - FieldPresence::Null - } else if metrics.as_object().map(|o| o.is_empty()).unwrap_or(true) { - FieldPresence::Missing + let metrics = if metrics.as_object().map(|o| o.is_empty()).unwrap_or(true) { + None } else { - FieldPresence::Value(metrics) - }; - - let turn_null = |field| { - story - .presence - .turn_nulls - .get(&turn.id) - .is_some_and(|fields| fields.contains(&field)) + Some(metrics) }; steps.push(AtifStep { step_id: turn.id, - timestamp: field_from_option( - turn.timestamp.clone(), - turn_null(StorylineTurnField::Timestamp), - ), + timestamp: turn.timestamp.clone(), source: turn.source.clone(), - model_name: field_from_option( - turn.model_name.clone(), - turn_null(StorylineTurnField::ModelName), - ), - reasoning_effort: field_from_option( - turn.reasoning_effort.clone(), - turn_null(StorylineTurnField::ReasoningEffort), - ), + model_name: turn.model_name.clone(), + reasoning_effort: turn.reasoning_effort.clone(), message: turn.message.clone(), - reasoning_content: field_from_option( - turn.reasoning_content.clone(), - turn_null(StorylineTurnField::ReasoningContent), - ), - tool_calls: field_from_option(tool_calls, turn_null(StorylineTurnField::ToolCalls)), - observation: field_from_option(observation, turn_null(StorylineTurnField::Observation)), + reasoning_content: turn.reasoning_content.clone(), + tool_calls, + observation, metrics, - extra: field_from_option(turn.extra.clone(), turn_null(StorylineTurnField::Extra)), - llm_call_count: field_from_option( - turn.llm_call_count, - turn_null(StorylineTurnField::LlmCallCount), - ), - is_copied_context: field_from_option( - turn.is_copied_context, - turn_null(StorylineTurnField::IsCopiedContext), - ), + extra: turn.extra.clone(), + llm_call_count: turn.llm_call_count, + is_copied_context: turn.is_copied_context, + unknown: Map::new(), }); } - let root_null = |field| story.presence.root_nulls.contains(&field); - let agent_null = |field| story.presence.agent_nulls.contains(&field); let subagent_trajectories = match embedded_children { - Some(children) => FieldPresence::Value(children), - None if root_null(StorylineRootField::SubagentTrajectories) => FieldPresence::Null, - None if story.child_session_ids.as_ref().is_some_and(Vec::is_empty) => { - FieldPresence::Value(Vec::new()) - } - None => FieldPresence::Missing, + Some(children) => Some(children), + None if story.child_session_ids.as_ref().is_some_and(Vec::is_empty) => Some(Vec::new()), + None => None, }; Ok(AtifTrajectory { @@ -476,54 +425,36 @@ fn storyline_to_atif_node( .schema_version .clone() .unwrap_or_else(|| "ATIF-v1.7".into()), - session_id: match story.presence.session_id { - PresenceState::Missing => FieldPresence::Missing, - PresenceState::Null => FieldPresence::Null, - PresenceState::Value => FieldPresence::Value(story.session_id.clone()), - }, - trajectory_id: field_from_option( - story.trajectory_id.clone(), - root_null(StorylineRootField::TrajectoryId), - ), + session_id: Some(story.session_id.clone()), + trajectory_id: story.trajectory_id.clone(), agent: AtifAgent { name: story .agent .name .clone() .unwrap_or_else(|| story.agent.id.clone()), - version: story.agent.version.clone().unwrap_or_default(), - model_name: field_from_option( - story.agent.model_name.clone(), - agent_null(StorylineAgentField::ModelName), - ), - tool_definitions: field_from_option( - story.agent.tool_definitions.clone(), - agent_null(StorylineAgentField::ToolDefinitions), - ), - extra: field_from_option( - story.agent.extra.clone(), - agent_null(StorylineAgentField::Extra), - ), + version: story + .agent + .version + .clone() + .unwrap_or_else(|| "unknown".into()), + model_name: story.agent.model_name.clone(), + tool_definitions: story.agent.tool_definitions.clone(), + extra: story.agent.extra.clone(), + unknown: Map::new(), }, steps, - notes: field_from_option(story.notes.clone(), root_null(StorylineRootField::Notes)), - final_metrics: field_from_option( - story.final_metrics.clone(), - root_null(StorylineRootField::FinalMetrics), - ), - continued_trajectory_ref: field_from_option( - story.continued_trajectory_ref.clone(), - root_null(StorylineRootField::ContinuedTrajectoryRef), - ), - extra: field_from_option(story.extra.clone(), root_null(StorylineRootField::Extra)), + notes: story.notes.clone(), + final_metrics: story.final_metrics.clone(), + continued_trajectory_ref: story.continued_trajectory_ref.clone(), + extra: story.extra.clone(), subagent_trajectories, + unknown: Map::new(), }) } /// Rebuild one or more ATIF trees from flattened Storyline documents. pub fn storylines_to_atif(stories: &[StorylineDocument]) -> Result> { - use std::collections::{HashMap, HashSet}; - fn key(story: &StorylineDocument) -> &str { story .trajectory_id @@ -601,14 +532,141 @@ pub fn storylines_to_atif(stories: &[StorylineDocument]) -> Result, + stories: &[StorylineDocument], +) -> Result> { + let single = documents.len() == 1; + let mut value = if documents.len() == 1 { + serde_json::to_value(&documents[0])? + } else { + serde_json::to_value(&documents)? + }; + + let indexes = stories + .iter() + .enumerate() + .map(|(index, story)| (story.document_id().to_string(), index)) + .collect::>(); + let referenced = stories + .iter() + .flat_map(|story| story.child_session_ids.iter().flatten().cloned()) + .collect::>(); + let root_indexes = stories + .iter() + .enumerate() + .filter_map(|(index, story)| (!referenced.contains(story.document_id())).then_some(index)) + .collect::>(); + + fn bind_tree( + story_index: usize, + pointer: &str, + stories: &[StorylineDocument], + indexes: &HashMap, + carriers: &mut Vec, + ) -> Result<()> { + carriers.push(CarrierBinding { + story_index, + pointer: pointer.to_string(), + }); + if let Some(children) = &stories[story_index].child_session_ids { + for (child_position, child) in children.iter().enumerate() { + let child_index = *indexes.get(child).ok_or_else(|| { + anyhow::anyhow!("Storyline child '{child}' has no matching document") + })?; + let child_pointer = pointer_join( + &pointer_join(pointer, "subagent_trajectories"), + &child_position.to_string(), + ); + bind_tree(child_index, &child_pointer, stories, indexes, carriers)?; + } + } + Ok(()) + } + + let mut carriers = Vec::new(); + for (root_position, story_index) in root_indexes.iter().copied().enumerate() { + let pointer = if root_indexes.len() == 1 { + String::new() + } else { + pointer_join("", &root_position.to_string()) + }; + bind_tree(story_index, &pointer, stories, &indexes, &mut carriers)?; + } + + let mut source_roots = BTreeMap::::new(); + for story_index in &root_indexes { + if let Some(source) = stories[*story_index].unknown_fields.sources.get("atif") { + if source_roots + .insert(source.source_document_id.clone(), *story_index) + .is_some() + { + anyhow::bail!( + "ATIF source document '{}' has multiple root trajectories", + source.source_document_id + ); + } + } + } + let carrier_by_story = carriers + .iter() + .map(|carrier| (carrier.story_index, carrier.pointer.clone())) + .collect::>(); + let mut merged = BTreeMap::>::new(); + for story in stories { + let Some(source) = story.unknown_fields.sources.get("atif") else { + continue; + }; + let fields = merged.entry(source.source_document_id.clone()).or_default(); + for (pointer, field_value) in &source.fields { + match fields.get(pointer) { + Some(existing) if existing != field_value => anyhow::bail!( + "ATIF source '{}' has conflicting unknown field at '{}'", + source.source_document_id, + pointer + ), + Some(_) => {} + None => { + fields.insert(pointer.clone(), field_value.clone()); + } + } + } + } + for (source_id, fields) in merged { + let root_story = source_roots.get(&source_id).ok_or_else(|| { + anyhow::anyhow!("ATIF source document '{source_id}' has no root trajectory") + })?; + let root_pointer = &carrier_by_story[root_story]; + let target = value + .pointer_mut(root_pointer) + .ok_or_else(|| anyhow::anyhow!("ATIF root carrier '{root_pointer}' is missing"))?; + for (pointer, field_value) in fields { + restore_json_pointer(target, &pointer, field_value, PointerWrite::InsertOnly) + .with_context(|| { + format!( + "restore ATIF unknown-field pointer '{pointer}' for trajectory '{}'", + stories[*root_story].document_id() + ) + })?; + } + } + + write_foreign_unknown_fields_envelope(DocumentFormat::Atif, &mut value, stories, &carriers)?; + if single { + Ok(vec![serde_json::from_value(value)?]) + } else { + Ok(serde_json::from_value(value)?) + } } #[cfg(test)] mod tests { use super::{atif_to_storyline, atif_to_storylines, storyline_to_atif, storylines_to_atif}; use crate::atif::AtifTrajectory; - use crate::FieldPresence; + use crate::StorylineDocument; #[test] fn malformed_atif_observation_is_not_silently_dropped() { @@ -638,7 +696,7 @@ mod tests { } #[test] - fn atif_tool_result_presence_round_trips_without_provenance() { + fn atif_tool_result_null_and_missing_canonicalize_to_absent() { let trajectory = AtifTrajectory::from_json_str( r#"{ "schema_version":"ATIF-v1.7", @@ -661,12 +719,9 @@ mod tests { let story = atif_to_storyline(&trajectory).unwrap(); assert_eq!(story.schema_version.as_deref(), Some("ATIF-v1.7")); let calls = story.turns[0].tool_calls.as_ref().unwrap(); - assert_eq!(calls[0].result, FieldPresence::Missing); - assert_eq!(calls[1].result, FieldPresence::Null); - assert_eq!( - calls[2].result, - FieldPresence::Value(serde_json::json!({"ok": true})) - ); + assert_eq!(calls[0].result, None); + assert_eq!(calls[1].result, None); + assert_eq!(calls[2].result, Some(serde_json::json!({"ok": true}))); assert!(calls.iter().all(|call| { !call .extra @@ -677,12 +732,12 @@ mod tests { let encoded = serde_json::to_value(storyline_to_atif(&story).unwrap()).unwrap(); let calls = encoded["steps"][0]["tool_calls"].as_array().unwrap(); assert!(calls[0].get("result").is_none()); - assert_eq!(calls[1]["result"], serde_json::Value::Null); + assert!(calls[1].get("result").is_none()); assert_eq!(calls[2]["result"], serde_json::json!({"ok": true})); } #[test] - fn atif_null_presence_and_trajectory_only_identity_round_trip() { + fn atif_null_fields_canonicalize_to_absent() { let input = serde_json::json!({ "schema_version": "ATIF-v1.7", "trajectory_id": "trajectory-only", @@ -714,7 +769,18 @@ mod tests { let trajectory = AtifTrajectory::from_json_str(&input.to_string()).unwrap(); let story = atif_to_storyline(&trajectory).unwrap(); let output = serde_json::to_value(storyline_to_atif(&story).unwrap()).unwrap(); - assert_eq!(output, input); + assert_eq!(output["session_id"], "trajectory-only"); + assert_eq!(output["trajectory_id"], "trajectory-only"); + assert!(output.get("notes").is_none()); + assert!(output["steps"][0].get("timestamp").is_none()); + } + + #[test] + fn cross_format_storyline_gets_a_valid_atif_agent_version() { + let story = StorylineDocument::new("session", "agent"); + let trajectory = storyline_to_atif(&story).unwrap(); + assert_eq!(trajectory.agent.version, "unknown"); + trajectory.validate().unwrap(); } #[test] @@ -751,7 +817,19 @@ mod tests { ); let rebuilt = storylines_to_atif(&stories).unwrap(); assert_eq!(rebuilt.len(), 1); - assert_eq!(serde_json::to_value(&rebuilt[0]).unwrap(), input); + let rebuilt = serde_json::to_value(&rebuilt[0]).unwrap(); + assert_eq!( + rebuilt["subagent_trajectories"][0]["session_id"], + "shared-run" + ); + assert_eq!( + rebuilt["subagent_trajectories"][0]["trajectory_id"], + "child-a" + ); + assert_eq!( + rebuilt["subagent_trajectories"][1]["trajectory_id"], + "child-b" + ); } #[test] diff --git a/crates/persisting-pchronicle/src/convert/events.rs b/crates/persisting-pchronicle/src/convert/events.rs index 8488d8ce..c1ca4c0c 100644 --- a/crates/persisting-pchronicle/src/convert/events.rs +++ b/crates/persisting-pchronicle/src/convert/events.rs @@ -309,7 +309,8 @@ fn events_to_storyline_unchecked(events: &[EventRecord]) -> Result = anyhow::Result; -use crate::atif::AtifTrajectory; use crate::convert::{atif_collection_to_storylines, storylines_to_actf, storylines_to_atif}; use crate::formats::actf::ActfDocument; +use crate::formats::unknown_fields::{ + attach_carried_unknown_fields, take_unknown_fields_envelope, validate_unknown_fields, + CarrierBinding, UnknownFieldLimits, +}; use crate::formats::{ has_openai_provenance, parse_openai_msg_corpus_value, recover_openai_msg_files, - synthesize_openai_msg_corpus, StorylineCollectionShape, StorylineDocument, + StorylineDocument, }; +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct DocumentCodecOptions { + pub unknown_fields: UnknownFieldLimits, +} + /// 将 AgenticMD 语义文档解码为权威 Storyline，不暴露 Markdown AST。 pub fn decode_agenticmd(input: &str) -> InputResult { crate::agenticmd::parse_agenticmd(input) @@ -41,38 +48,82 @@ pub fn decode_json_storylines( input: &str, relative_path: impl AsRef, ) -> InputResult> { - match format { + decode_json_storylines_with_options( + format, + input, + relative_path, + DocumentCodecOptions::default(), + ) +} + +pub fn decode_json_storylines_with_options( + format: DocumentFormat, + input: &str, + relative_path: impl AsRef, + options: DocumentCodecOptions, +) -> InputResult> { + options.unknown_fields.validate()?; + let stories = match format { DocumentFormat::Atif => { let value: serde_json::Value = serde_json::from_str(input) .map_err(|error| InputIssue::invalid(error.to_string()))?; - let (values, array) = match value { - serde_json::Value::Array(values) => (values, true), - value => (vec![value], false), + let values = match value { + serde_json::Value::Array(values) => values, + value => vec![value], }; if values.is_empty() { return Err(InputIssue::unsupported("ATIF document cannot be empty")); } let mut stories = Vec::new(); - for (ordinal, value) in values.into_iter().enumerate() { - let trajectory = AtifTrajectory::from_json_str(&value.to_string())?; - let ordinal = i64::try_from(ordinal) - .map_err(|_| InputIssue::invalid("ATIF collection ordinal overflow"))?; - let shape = if array { - StorylineCollectionShape::Sequence - } else { - StorylineCollectionShape::Single - }; + for value in values { stories.extend( - atif_collection_to_storylines(&trajectory, shape, ordinal) + atif_collection_to_storylines(value) .map_err(|error| InputIssue::invalid(error.to_string()))?, ); } Ok(stories) } DocumentFormat::Actf => { - let document = ActfDocument::from_json_str(input)?; - crate::convert::actf_to_storylines(&document) - .map_err(|error| InputIssue::invalid(error.to_string())) + let mut value: serde_json::Value = serde_json::from_str(input) + .map_err(|error| InputIssue::invalid(error.to_string()))?; + let envelope = take_unknown_fields_envelope(&mut value)?; + let document: ActfDocument = serde_json::from_value(value) + .map_err(|error| InputIssue::invalid(error.to_string()))?; + document.validate()?; + let mut stories = crate::convert::actf_to_storylines(&document) + .map_err(|error| InputIssue::invalid(error.to_string()))?; + let owned_counts = stories + .iter() + .map(|story| story.unknown_key_counts.get("actf").cloned()) + .collect::>(); + let carriers = stories + .iter() + .enumerate() + .map(|(story_index, story)| CarrierBinding { + story_index, + pointer: format!( + "/attempts/{}", + story + .attempt_id + .as_deref() + .unwrap_or("1") + .replace('~', "~0") + .replace('/', "~1") + ), + }) + .collect::>(); + attach_carried_unknown_fields( + envelope, + &carriers, + &mut stories, + options.unknown_fields, + )?; + for (story, owned) in stories.iter_mut().zip(owned_counts) { + if let Some(owned) = owned { + story.unknown_key_counts.insert("actf".into(), owned); + } + } + Ok(stories) } DocumentFormat::OpenaiMsg => { let value = serde_json::from_str(input) @@ -82,7 +133,11 @@ pub fn decode_json_storylines( unsupported => Err(InputIssue::unsupported(format!( "'{unsupported}' is not a peripheral JSON document format" ))), + }?; + for story in &stories { + validate_unknown_fields(&story.unknown_fields, options.unknown_fields)?; } + Ok(stories) } /// 将权威 Storyline 编码为一个外围 JSON 文档。 @@ -90,13 +145,22 @@ pub fn encode_json_storylines( format: DocumentFormat, stories: &[StorylineDocument], ) -> Result { + encode_json_storylines_with_options(format, stories, DocumentCodecOptions::default()) +} + +pub fn encode_json_storylines_with_options( + format: DocumentFormat, + stories: &[StorylineDocument], + options: DocumentCodecOptions, +) -> Result { + options.unknown_fields.validate()?; + for story in stories { + validate_unknown_fields(&story.unknown_fields, options.unknown_fields)?; + } match format { DocumentFormat::Atif => { - let (stories, collection_shape) = prepare_atif_collection(stories)?; - let documents = storylines_to_atif(&stories)?; - if collection_shape == Some(StorylineCollectionShape::Sequence) { - Ok(serde_json::to_value(documents)?) - } else if documents.len() == 1 { + let documents = storylines_to_atif(stories)?; + if documents.len() == 1 { Ok(serde_json::to_value(&documents[0])?) } else { Ok(serde_json::to_value(documents)?) @@ -108,87 +172,6 @@ pub fn encode_json_storylines( } } -fn prepare_atif_collection( - stories: &[StorylineDocument], -) -> Result<(Vec, Option)> { - let mut shape = None; - let mut missing_shape = false; - let mut has_ordinal = false; - let mut missing_ordinal = false; - for story in stories { - story.validate()?; - match story.presence.collection_shape { - Some(value) if shape.is_some_and(|current| current != value) => { - anyhow::bail!("Storyline collection contains conflicting container shapes"); - } - Some(value) => shape = Some(value), - None => missing_shape = true, - } - if story.presence.collection_ordinal.is_some() { - has_ordinal = true; - } else { - missing_ordinal = true; - } - } - if shape.is_some() && missing_shape { - anyhow::bail!("Storyline collection mixes declared and undeclared container shapes"); - } - if has_ordinal && missing_ordinal { - anyhow::bail!("Storyline collection mixes declared and undeclared ordinals"); - } - - let mut ordered = stories.to_vec(); - if has_ordinal { - ordered.sort_by_key(|story| story.presence.collection_ordinal); - } - - let by_id = ordered - .iter() - .map(|story| (story.document_id().to_string(), story)) - .collect::>(); - if by_id.len() != ordered.len() { - anyhow::bail!("Storyline collection contains duplicate document identities"); - } - let referenced = ordered - .iter() - .flat_map(|story| story.child_session_ids.iter().flatten().cloned()) - .collect::>(); - let roots = ordered - .iter() - .filter(|story| !referenced.contains(story.document_id())) - .collect::>(); - let mut root_ordinals = HashSet::new(); - for root in &roots { - if let Some(ordinal) = root.presence.collection_ordinal { - if !root_ordinals.insert(ordinal) { - anyhow::bail!("duplicate Storyline root collection ordinal {ordinal}"); - } - } - } - for story in &ordered { - if let Some(children) = &story.child_session_ids { - for child_id in children { - let child = by_id.get(child_id).ok_or_else(|| { - anyhow::anyhow!("Storyline child '{child_id}' has no matching document") - })?; - if child.presence.collection_ordinal != story.presence.collection_ordinal { - anyhow::bail!( - "Storyline child '{child_id}' has a different collection ordinal" - ); - } - } - } - } - if shape == Some(StorylineCollectionShape::Single) - && (roots.len() != 1 || root_ordinals.iter().any(|ordinal| *ordinal != 0)) - { - anyhow::bail!( - "single-document Storyline collection must have exactly one root at ordinal zero" - ); - } - Ok((ordered, shape)) -} - fn encode_openai_storylines(stories: &[StorylineDocument]) -> Result { if stories.iter().any(has_openai_provenance) { let mut files = recover_openai_msg_files(stories)?; @@ -200,7 +183,7 @@ fn encode_openai_storylines(stories: &[StorylineDocument]) -> Result serde_json::Value { + serde_json::json!({ "schema_version": "ATIF-v1.7", "trajectory_id": "one", "agent": {"name": "agent", "version": "1"}, "steps": [] - }]); + }) + } + + #[test] + fn atif_singleton_object_and_array_encode_canonically() { + let object = atif_fixture_value(); + let from_object = + decode_json_storylines(DocumentFormat::Atif, &object.to_string(), "a.json").unwrap(); + let from_array = decode_json_storylines( + DocumentFormat::Atif, + &serde_json::json!([object]).to_string(), + "a.json", + ) + .unwrap(); + assert_eq!( + encode_json_storylines(DocumentFormat::Atif, &from_object).unwrap(), + encode_json_storylines(DocumentFormat::Atif, &from_array).unwrap(), + ); + } + + #[test] + fn atif_unknown_fields_round_trip_canonically() { + let input = serde_json::json!({ + "schema_version": "ATIF-v1.7", + "session_id": null, + "trajectory_id": "t1", + "vendor_root": null, + "agent": {"name": "a", "version": "1", "vendor_agent": {"x": 1}}, + "steps": [{ + "step_id": 1, "source": "user", "message": "hi", + "vendor_step": [1, 2], "0": "numeric-object-key" + }] + }); let stories = - decode_json_storylines(DocumentFormat::Atif, &input.to_string(), "one.json").unwrap(); + decode_json_storylines(DocumentFormat::Atif, &input.to_string(), "t.json").unwrap(); assert_eq!( - encode_json_storylines(DocumentFormat::Atif, &stories).unwrap(), - input + stories[0].unknown_fields.sources["atif"].fields["/vendor_root"], + serde_json::Value::Null ); + assert_eq!( + stories[0].unknown_key_counts["atif"]["/steps/*/vendor_step"], + 1 + ); + assert_eq!(stories[0].unknown_key_counts["atif"]["/steps/*/0"], 1); + let output = encode_json_storylines(DocumentFormat::Atif, &stories).unwrap(); + assert_eq!(output["vendor_root"], serde_json::Value::Null); + assert!(output.get("session_id").is_some()); } #[test] @@ -318,21 +341,24 @@ mod tests { } #[test] - fn atif_encoder_rejects_conflicting_collection_metadata() { - let mut first = StorylineDocument::new("first", "agent"); - first.presence.collection_shape = Some(StorylineCollectionShape::Sequence); - first.presence.collection_ordinal = Some(0); - let mut second = StorylineDocument::new("second", "agent"); - second.presence.collection_shape = Some(StorylineCollectionShape::Single); - second.presence.collection_ordinal = Some(0); - assert!(encode_json_storylines(DocumentFormat::Atif, &[first, second]).is_err()); + fn unknown_fields_options_are_validated_before_json_decode() { + let options = DocumentCodecOptions { + unknown_fields: UnknownFieldLimits { + max_fields: 0, + max_bytes: 1, + }, + }; + let decode_error = decode_json_storylines_with_options( + DocumentFormat::Atif, + "not JSON", + "invalid.json", + options, + ) + .unwrap_err(); + assert!(decode_error.to_string().contains("unknown field limit")); - let mut first = StorylineDocument::new("first", "agent"); - first.presence.collection_shape = Some(StorylineCollectionShape::Sequence); - first.presence.collection_ordinal = Some(0); - let mut second = StorylineDocument::new("second", "agent"); - second.presence.collection_shape = Some(StorylineCollectionShape::Sequence); - second.presence.collection_ordinal = Some(0); - assert!(encode_json_storylines(DocumentFormat::Atif, &[first, second]).is_err()); + let encode_error = + encode_json_storylines_with_options(DocumentFormat::Atif, &[], options).unwrap_err(); + assert!(encode_error.to_string().contains("unknown field limit")); } } diff --git a/crates/persisting-pchronicle/src/formats/mod.rs b/crates/persisting-pchronicle/src/formats/mod.rs index 89c2c016..59d69835 100644 --- a/crates/persisting-pchronicle/src/formats/mod.rs +++ b/crates/persisting-pchronicle/src/formats/mod.rs @@ -1,4 +1,4 @@ -//! Codecs for each [`crate::DocumentFormat`]. +//! Codecs for each [`crate::document::DocumentFormat`]. pub mod actf; pub mod detect; @@ -6,11 +6,10 @@ pub mod events; pub mod llm; pub mod openai_corpus; pub mod storyline; +pub mod unknown_fields; pub use detect::detect_format; pub use events::{EventIdentity, EventRecord}; pub(crate) use openai_corpus::has_openai_provenance; -pub use openai_corpus::{ - parse_openai_msg_corpus_value, recover_openai_msg_files, synthesize_openai_msg_corpus, -}; -pub use storyline::{StorylineCollectionShape, StorylineDocument}; +pub use openai_corpus::{parse_openai_msg_corpus_value, recover_openai_msg_files}; +pub use storyline::StorylineDocument; diff --git a/crates/persisting-pchronicle/src/formats/openai_corpus.rs b/crates/persisting-pchronicle/src/formats/openai_corpus.rs index 8323a147..2a03b1d7 100644 --- a/crates/persisting-pchronicle/src/formats/openai_corpus.rs +++ b/crates/persisting-pchronicle/src/formats/openai_corpus.rs @@ -2,21 +2,28 @@ //! //! The reader accepts either a top-level row array or a `session_steps` //! envelope containing rows from many sessions. Unmapped container and row -//! fields are retained as controlled residuals so strict recovery can rebuild +//! fields are retained as controlled unknown fields so strict recovery can rebuild //! the original source document. use std::collections::{HashMap, HashSet}; use std::path::{Component, Path, PathBuf}; +use anyhow::Context as _; use chrono::{SecondsFormat, TimeZone, Utc}; use serde_json::{json, Map, Value}; +use crate::format::DocumentFormat; use crate::formats::storyline::{ StorylineAgent, StorylineDocument, StorylineToolCall, StorylineTurn, }; +use crate::formats::unknown_fields::{ + attach_carried_unknown_fields, normalize_openai_pointer, restore_json_pointer, + take_unknown_fields_envelope, validate_unknown_fields_with, + write_foreign_unknown_fields_envelope, CarrierBinding, PointerWrite, SourceUnknownFields, + UnknownFieldLimits, +}; use crate::{InputIssue, InputResult, Result}; -const OPENAI_EXTENSION_KEY: &str = "persisting.dev/openai-msg/v1"; const ROW_METRIC_FIELDS: &[&str] = &[ "reward", "step_reward", @@ -41,8 +48,10 @@ pub fn parse_openai_msg_corpus_value( let relative_path = validate_input_relative_path(relative_path.as_ref())? .to_string_lossy() .into_owned(); - let (kind, envelope, records) = match document { - Value::Array(records) => ("array", None, records.clone()), + let mut document = document.clone(); + let carried_envelope = take_unknown_fields_envelope(&mut document)?; + let (records, root_unknown) = match &document { + Value::Array(records) => (records.clone(), Map::new()), Value::Object(root) => { let records = root .get("session_steps") @@ -53,7 +62,7 @@ pub fn parse_openai_msg_corpus_value( .clone(); let mut metadata = root.clone(); metadata.remove("session_steps"); - ("envelope", Some(Value::Object(metadata)), records) + (records, metadata) } _ => { return Err(InputIssue::invalid( @@ -62,11 +71,6 @@ pub fn parse_openai_msg_corpus_value( } }; - let file_metadata = json!({ - "relative_path": relative_path, - "document_kind": kind, - "envelope": envelope, - }); let mut groups: Vec<(String, Vec<(usize, Value)>)> = Vec::new(); let mut group_indexes = HashMap::::new(); for (ordinal, record) in records.into_iter().enumerate() { @@ -91,164 +95,252 @@ pub fn parse_openai_msg_corpus_value( return Err(InputIssue::unsupported("OpenAI corpus cannot be empty")); } - groups - .into_iter() - .map(|(session_id, records)| { - rows_to_storyline(&session_id, records, &relative_path, &file_metadata) - }) - .collect() -} - -/// Recover original OpenAI files from Storylines produced by the corpus reader. -/// -/// This is intentionally strict: Storylines without complete lossless metadata -/// are rejected instead of being silently synthesized from normalized fields. -pub fn recover_openai_msg_files( - stories: &[StorylineDocument], -) -> Result> { - #[derive(Clone)] - struct FileGroup { - kind: String, - envelope: Option, - records: Vec<(u64, Value)>, - } - - let mut files = HashMap::::new(); - for story in stories { - let file = story - .extra - .as_ref() - .and_then(|extra| extra.get(OPENAI_EXTENSION_KEY)) - .and_then(Value::as_object) - .ok_or_else(|| { - anyhow::anyhow!( - "Storyline '{}' has no lossless OpenAI file metadata", - story.session_id - ) - })?; - let relative_path = file - .get("relative_path") - .and_then(Value::as_str) - .ok_or_else(|| anyhow::anyhow!("OpenAI file metadata missing relative_path"))?; - let relative_path = validate_relative_path(Path::new(relative_path))?; - let kind = file - .get("document_kind") - .and_then(Value::as_str) - .filter(|kind| matches!(*kind, "array" | "envelope")) - .ok_or_else(|| anyhow::anyhow!("invalid OpenAI document_kind"))? - .to_string(); - let envelope = file.get("envelope").filter(|v| !v.is_null()).cloned(); - - let group = files - .entry(relative_path.clone()) - .or_insert_with(|| FileGroup { - kind: kind.clone(), - envelope: envelope.clone(), - records: Vec::new(), + let mut stories = Vec::with_capacity(groups.len()); + let mut carriers = Vec::new(); + for (session_id, records) in groups { + let story_index = stories.len(); + let mut story = rows_to_storyline(&session_id, records.clone(), &relative_path)?; + capture_openai_unknowns(&mut story, &relative_path, &root_unknown, &records)?; + story.unknown_key_counts = validate_unknown_fields_with( + &story.unknown_fields, + UnknownFieldLimits::default(), + normalize_openai_pointer, + )?; + for (ordinal, _) in records { + carriers.push(CarrierBinding { + story_index, + pointer: format!("/session_steps/{ordinal}"), }); - if group.kind != kind || group.envelope != envelope { - anyhow::bail!( - "conflicting OpenAI file metadata for {}", - relative_path.display() - ); } + stories.push(story); + } + let owned_counts = stories + .iter() + .map(|story| story.unknown_key_counts.get("openai-msg").cloned()) + .collect::>(); + attach_carried_unknown_fields( + carried_envelope, + &carriers, + &mut stories, + UnknownFieldLimits::default(), + )?; + for (story, owned) in stories.iter_mut().zip(owned_counts) { + if let Some(owned) = owned { + story.unknown_key_counts.insert("openai-msg".into(), owned); + } + } + Ok(stories) +} - for turn in &story.turns { - if turn.source == "user" { - continue; +fn capture_openai_unknowns( + story: &mut StorylineDocument, + source_document_id: &str, + root_unknown: &Map, + records: &[(usize, Value)], +) -> InputResult<()> { + story + .unknown_fields + .sources + .entry("openai-msg".into()) + .or_insert_with(|| SourceUnknownFields { + source_document_id: source_document_id.to_string(), + fields: Default::default(), + }); + insert_openai_map(story, source_document_id, "", root_unknown)?; + for (ordinal, record) in records { + let row = record.as_object().expect("rows were validated as objects"); + let row_prefix = format!("/session_steps/{ordinal}"); + for (key, value) in row { + if !is_canonical_openai_row_key(key) { + story.unknown_fields.insert( + "openai-msg", + source_document_id, + pointer_join(&row_prefix, key), + value.clone(), + )?; } - let extra = turn.extra.as_ref().ok_or_else(|| { - anyhow::anyhow!( - "Storyline '{}' step {} has no OpenAI provenance", - story.session_id, - turn.id - ) - })?; - let Some(record) = extra.get(OPENAI_EXTENSION_KEY).and_then(Value::as_object) else { - anyhow::bail!( - "Storyline '{}' step {} has no OpenAI residual", - story.session_id, - turn.id - ); + } + if let Some(messages) = row.get("messages").and_then(Value::as_array) { + let output_index = match select_output_message(row).map(|(_, location)| location) { + Some(OutputLocation::Message(index)) => Some(index), + _ => None, }; - let record_path = record - .get("relative_path") - .and_then(Value::as_str) - .ok_or_else(|| anyhow::anyhow!("OpenAI record missing relative_path"))?; - if validate_relative_path(Path::new(record_path))? != relative_path { - anyhow::bail!( - "OpenAI record path conflicts with Storyline '{}' file metadata", - story.session_id - ); + let mut request_index = 0usize; + for (index, message) in messages.iter().enumerate() { + let Some(message) = message.as_object() else { + continue; + }; + let prefix = if output_index == Some(index) { + pointer_join(&row_prefix, "response") + } else { + let prefix = pointer_join( + &pointer_join(&row_prefix, "messages"), + &request_index.to_string(), + ); + request_index += 1; + prefix + }; + capture_openai_message(story, source_document_id, &prefix, message)?; } - let ordinal = record - .get("ordinal") - .and_then(Value::as_u64) - .ok_or_else(|| anyhow::anyhow!("OpenAI record missing ordinal"))?; - let raw = recover_record(story, turn, record)?; - group.records.push((ordinal, raw)); + } + if let Some(response) = row.get("response").and_then(Value::as_object) { + capture_openai_message( + story, + source_document_id, + &pointer_join(&row_prefix, "response"), + response, + )?; } } + Ok(()) +} - let mut output = Vec::with_capacity(files.len()); - for (relative_path, mut group) in files { - group.records.sort_by_key(|(ordinal, _)| *ordinal); - for pair in group.records.windows(2) { - if pair[0].0 == pair[1].0 { - anyhow::bail!( - "duplicate OpenAI row ordinal {} in {}", - pair[0].0, - relative_path.display() - ); - } +fn is_canonical_openai_row_key(key: &str) -> bool { + matches!( + key, + "session_id" + | "step_id" + | "id" + | "messages" + | "response" + | "agent_model" + | "llm_model" + | "run_id" + | "run_bucket" + | "job_id" + | "created_at" + | "meta" + | "env_state" + | "metrics" + | "call_id" + | "agent_id" + | "group_id" + | "env_name" + ) || ROW_METRIC_FIELDS.contains(&key) +} + +fn capture_openai_message( + story: &mut StorylineDocument, + source_document_id: &str, + prefix: &str, + message: &Map, +) -> InputResult<()> { + for (key, value) in message { + if !matches!( + key.as_str(), + "role" | "content" | "name" | "tool_call_id" | "tool_calls" + ) { + story.unknown_fields.insert( + "openai-msg", + source_document_id, + pointer_join(prefix, key), + value.clone(), + )?; } - // Ordinals are ordering keys, not a completeness proof. Callers may - // intentionally export a filtered set of complete trajectories from - // one source file, so gaps are valid while duplicates are not. - let records = group - .records - .into_iter() - .map(|(_, record)| record) - .collect::>(); - let document = match group.kind.as_str() { - "array" => Value::Array(records), - "envelope" => { - let mut envelope = group - .envelope - .and_then(|value| value.as_object().cloned()) - .ok_or_else(|| { - anyhow::anyhow!( - "OpenAI envelope metadata missing for {}", - relative_path.display() - ) - })?; - envelope.insert("session_steps".into(), Value::Array(records)); - Value::Object(envelope) + } + if let Some(calls) = message.get("tool_calls").and_then(Value::as_array) { + for (index, call) in calls.iter().enumerate() { + let Some(call) = call.as_object() else { + continue; + }; + let call_prefix = pointer_join(&pointer_join(prefix, "tool_calls"), &index.to_string()); + for (key, value) in call { + if !matches!(key.as_str(), "id" | "type" | "function") { + story.unknown_fields.insert( + "openai-msg", + source_document_id, + pointer_join(&call_prefix, key), + value.clone(), + )?; + } } - kind => { - anyhow::bail!( - "invalid OpenAI document kind '{}' while recovering {}", - kind, - relative_path.display() - ) + if let Some(function) = call.get("function").and_then(Value::as_object) { + for (key, value) in function { + if !matches!(key.as_str(), "name" | "arguments") { + story.unknown_fields.insert( + "openai-msg", + source_document_id, + pointer_join(&pointer_join(&call_prefix, "function"), key), + value.clone(), + )?; + } + } } - }; - output.push(RecoveredOpenaiMsgFile { - relative_path, - document, - }); + } + } + Ok(()) +} + +fn insert_openai_map( + story: &mut StorylineDocument, + source_document_id: &str, + prefix: &str, + fields: &Map, +) -> InputResult<()> { + for (key, value) in fields { + story.unknown_fields.insert( + "openai-msg", + source_document_id, + pointer_join(prefix, key), + value.clone(), + )?; } + Ok(()) +} + +fn pointer_join(parent: &str, token: &str) -> String { + format!("{parent}/{}", token.replace('~', "~0").replace('/', "~1")) +} + +/// Recover original OpenAI files from Storylines produced by the corpus reader. +/// +/// This is intentionally strict: Storylines without complete lossless metadata +/// are rejected instead of being silently synthesized from normalized fields. +pub fn recover_openai_msg_files( + stories: &[StorylineDocument], +) -> Result> { + let mut groups = HashMap::>::new(); + for story in stories { + let source_id = story + .unknown_fields + .sources + .get("openai-msg") + .map(|source| source.source_document_id.as_str()) + .or_else(|| { + story + .extra + .as_ref() + .and_then(|extra| extra.get("openai_source_document_id")) + .and_then(Value::as_str) + }) + .ok_or_else(|| { + anyhow::anyhow!("cannot mix OpenAI unknown fields and unrelated Storylines") + })?; + groups + .entry(source_id.to_string()) + .or_default() + .push(story.clone()); + } + let mut output = groups + .into_iter() + .map(|(relative_path, group)| { + Ok(RecoveredOpenaiMsgFile { + relative_path: validate_relative_path(Path::new(&relative_path))?, + document: storylines_to_openai_value(&group)?, + }) + }) + .collect::>>()?; output.sort_by(|left, right| left.relative_path.cmp(&right.relative_path)); Ok(output) } pub(crate) fn has_openai_provenance(story: &StorylineDocument) -> bool { - story - .extra - .as_ref() - .and_then(|extra| extra.get(OPENAI_EXTENSION_KEY)) - .and_then(Value::as_object) - .is_some() + story.unknown_fields.sources.contains_key("openai-msg") + || story + .extra + .as_ref() + .and_then(|extra| extra.get("openai_source_document_id")) + .is_some() } /// Explicitly synthesize an OpenAI message row array from Storyline semantics. @@ -256,10 +348,13 @@ pub(crate) fn has_openai_provenance(story: &StorylineDocument) -> bool { /// This is a cross-format projection, not a lossless recovery operation. Use /// [`recover_openai_msg_files`] when the Storylines originated from an OpenAI /// corpus and exact JSON-model recovery is required. -pub fn synthesize_openai_msg_corpus(stories: &[StorylineDocument]) -> Result { - let mut records = Vec::new(); - for story in stories { - story.validate()?; +pub(crate) fn synthesize_openai_msg_corpus(stories: &[StorylineDocument]) -> Result { + let mut records = Vec::<(Option, usize, usize, Value)>::new(); + let mut sequence = 0usize; + for (story_index, story) in stories.iter().enumerate() { + if story.session_id.is_empty() || story.agent.id.is_empty() { + anyhow::bail!("invalid Storyline identity for OpenAI conversion"); + } let mut index = 0usize; while index < story.turns.len() { let turn = &story.turns[index]; @@ -295,60 +390,213 @@ pub fn synthesize_openai_msg_corpus(stories: &[StorylineDocument]) -> Result(response) }) - }); + .transpose()?; let call_id = agent .and_then(|turn| turn.extra.as_ref()) .and_then(|extra| extra.get("call_id")) .and_then(Value::as_str) .unwrap_or(""); - records.push(json!({ - "id": format!("step-{}", output.id), - "session_id": story.session_id, - "step_id": output.id, - "job_id": "", - "agent_id": story.agent.id, - "group_id": "", - "env_name": "", - "llm_model": agent.and_then(|turn| turn.model_name.clone()).unwrap_or_default(), - "step_reward": 0.0, - "reward": 0.0, - "is_terminal": index >= story.turns.len(), - "is_truncated": false, - "is_session_completed": index >= story.turns.len(), - "is_trainable": true, - "created_at": output.timestamp.clone().unwrap_or_default(), - "messages": messages, - "response": response, - "run_bucket": story.run_id.clone().unwrap_or_default(), - "call_id": call_id, - })); + let ordinal = agent + .and_then(|turn| turn.extra.as_ref()) + .and_then(|extra| extra.get("openai_source_ordinal")) + .and_then(Value::as_u64); + let step_id = agent + .and_then(|turn| turn.extra.as_ref()) + .and_then(|extra| extra.get("openai_step_id")) + .and_then(Value::as_i64) + .unwrap_or(output.id); + records.push(( + ordinal, + sequence, + story_index, + json!({ + "id": call_id, + "session_id": story.session_id, + "step_id": step_id, + "job_id": "", + "agent_id": story.agent.id, + "group_id": "", + "env_name": "", + "llm_model": agent.and_then(|turn| turn.model_name.clone()).unwrap_or_default(), + "step_reward": 0.0, + "reward": 0.0, + "is_terminal": index >= story.turns.len(), + "is_truncated": false, + "is_session_completed": index >= story.turns.len(), + "is_trainable": true, + "created_at": output.timestamp.clone(), + "messages": messages, + "response": response, + "run_bucket": story.run_id.clone().unwrap_or_default(), + "call_id": call_id, + }), + )); + sequence += 1; } } - Ok(Value::Array(records)) + records.sort_by_key(|(ordinal, sequence, _, _)| { + (ordinal.is_none(), ordinal.unwrap_or(u64::MAX), *sequence) + }); + Ok(json!({ + "session_steps": records.into_iter().map(|(_, _, _, row)| row).collect::>() + })) +} + +pub(crate) fn storylines_to_openai_value(stories: &[StorylineDocument]) -> Result { + let mut value = synthesize_openai_msg_corpus(stories)?; + let rows = value["session_steps"] + .as_array() + .ok_or_else(|| anyhow::anyhow!("canonical OpenAI envelope lost session_steps"))?; + let story_by_session = stories + .iter() + .enumerate() + .map(|(index, story)| (story.session_id.as_str(), index)) + .collect::>(); + let mut hint_by_row = HashMap::<(String, i64), (usize, usize)>::new(); + for (story_index, story) in stories.iter().enumerate() { + for turn in story.turns.iter().filter(|turn| turn.source == "agent") { + let Some(extra) = turn.extra.as_ref() else { + continue; + }; + let Some(ordinal) = extra.get("openai_source_ordinal").and_then(Value::as_u64) else { + continue; + }; + let step_id = extra + .get("openai_step_id") + .and_then(Value::as_i64) + .unwrap_or(turn.id); + let ordinal = usize::try_from(ordinal).context("OpenAI source ordinal overflow")?; + hint_by_row.insert((story.session_id.clone(), step_id), (ordinal, story_index)); + } + } + let mut ordinal_to_target = HashMap::::new(); + let mut carriers = Vec::new(); + for (target_index, row) in rows.iter().enumerate() { + let session = row["session_id"] + .as_str() + .ok_or_else(|| anyhow::anyhow!("canonical OpenAI row missing session_id"))?; + let step_id = row["step_id"] + .as_i64() + .ok_or_else(|| anyhow::anyhow!("canonical OpenAI row missing step_id"))?; + let story_index = *story_by_session.get(session).ok_or_else(|| { + anyhow::anyhow!("canonical OpenAI row has unknown session '{session}'") + })?; + if let Some((ordinal, hinted_story)) = hint_by_row.get(&(session.to_string(), step_id)) { + if *hinted_story != story_index + || ordinal_to_target.insert(*ordinal, target_index).is_some() + { + anyhow::bail!("duplicate OpenAI source row carrier for ordinal {ordinal}"); + } + } + carriers.push(CarrierBinding { + story_index, + pointer: format!("/session_steps/{target_index}"), + }); + } + + let mut merged = std::collections::BTreeMap::::new(); + let mut source_id = None::; + for story in stories { + let Some(source) = story.unknown_fields.sources.get("openai-msg") else { + continue; + }; + if source_id + .as_ref() + .is_some_and(|id| id != &source.source_document_id) + { + anyhow::bail!("one OpenAI output cannot merge multiple source documents"); + } + source_id = Some(source.source_document_id.clone()); + for (pointer, field_value) in &source.fields { + let target_pointer = remap_openai_pointer(pointer, &ordinal_to_target)?; + match merged.get(&target_pointer) { + Some(existing) if existing != field_value => { + anyhow::bail!("OpenAI unknown-field conflict at '{target_pointer}'") + } + Some(_) => {} + None => { + merged.insert(target_pointer, field_value.clone()); + } + } + } + } + for (pointer, field_value) in merged { + restore_json_pointer(&mut value, &pointer, field_value, PointerWrite::InsertOnly) + .with_context(|| format!("restore OpenAI unknown field '{pointer}'"))?; + } + write_foreign_unknown_fields_envelope( + DocumentFormat::OpenaiMsg, + &mut value, + stories, + &carriers, + )?; + Ok(value) +} + +fn remap_openai_pointer( + pointer: &str, + ordinal_to_target: &HashMap, +) -> Result { + if !pointer.starts_with("/session_steps/") { + return Ok(pointer.to_string()); + } + let suffix = &pointer["/session_steps/".len()..]; + let (ordinal, rest) = suffix.split_once('/').unwrap_or((suffix, "")); + let ordinal = ordinal + .parse::() + .with_context(|| format!("invalid OpenAI source ordinal in '{pointer}'"))?; + let target = ordinal_to_target.get(&ordinal).ok_or_else(|| { + anyhow::anyhow!("OpenAI unknown field references filtered or missing source row {ordinal}") + })?; + Ok(if rest.is_empty() { + format!("/session_steps/{target}") + } else { + format!("/session_steps/{target}/{rest}") + }) } fn rows_to_storyline( session_id: &str, mut records: Vec<(usize, Value)>, relative_path: &str, - file_metadata: &Value, ) -> InputResult { records.sort_by_key(|(_, row)| row.get("step_id").and_then(Value::as_i64)); let mut seen_steps = HashSet::new(); let mut turns = Vec::with_capacity(records.len().saturating_mul(2)); let mut agent_source = None; + let mut first_agent_id = None; let mut first_model: Option = None; let mut run_id: Option = None; let mut next_turn_id = 1_i64; @@ -377,6 +625,13 @@ fn rows_to_storyline( if first_model.is_none() { first_model = model.clone(); } + if first_agent_id.is_none() { + first_agent_id = row + .get("agent_id") + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + .map(str::to_string); + } if agent_source.is_none() { agent_source = meta .as_ref() @@ -431,12 +686,16 @@ fn rows_to_storyline( let call_id = row .get("id") .and_then(Value::as_str) - .filter(|value| !value.is_empty()) .map(str::to_string) .unwrap_or_else(|| format!("step-{step_id}")); let request_messages = row.get("messages").cloned(); let user_message = last_user_message(request_messages.as_ref()); - let user_turn_id = user_message.as_ref().map(|_| next_turn_id); + let observation = parse_tool_results(request_messages.as_ref()); + let request_messages = request_message_context( + request_messages, + user_message.as_ref().map(|(index, _)| *index), + output_location, + ); if let Some((_, message)) = user_message.as_ref() { turns.push(StorylineTurn { id: next_turn_id, @@ -454,13 +713,7 @@ fn rows_to_storyline( is_copied_context: None, latency_ms: None, ttft_ms: None, - extra: Some(json!({ - "call_id": call_id, - OPENAI_EXTENSION_KEY: { - "kind": "request", - "openai_step_id": step_id, - } - })), + extra: Some(json!({"call_id": call_id})), }); next_turn_id += 1; } @@ -478,7 +731,7 @@ fn rows_to_storyline( reasoning_content: None, reasoning_effort: None, tool_calls, - observation: None, + observation, metrics, model_name: model, llm_call_count: Some(1), @@ -487,23 +740,17 @@ fn rows_to_storyline( ttft_ms, extra: Some(json!({ "call_id": call_id, - OPENAI_EXTENSION_KEY: record_residual( - row, - relative_path, - ordinal, - step_id, - user_message.as_ref().map(|(index, _)| *index), - user_turn_id, - output_location, - env_state.as_ref(), - ) + "openai_source_ordinal": ordinal, + "openai_step_id": step_id, + "request_messages": request_messages, })), }); next_turn_id += 1; } let final_metrics = turns.last().and_then(|turn| turn.metrics.clone()); - let agent_id = agent_source + let agent_id = first_agent_id + .or(agent_source) .or_else(|| first_model.clone()) .unwrap_or_else(|| "openai-import".into()); Ok(StorylineDocument { @@ -525,322 +772,101 @@ fn rows_to_storyline( notes: None, final_metrics, continued_trajectory_ref: None, - extra: Some(json!({ OPENAI_EXTENSION_KEY: file_metadata })), - presence: Default::default(), + extra: Some(json!({"openai_source_document_id": relative_path})), + unknown_fields: Default::default(), + unknown_key_counts: Default::default(), turns, }) } -fn last_user_message(messages: Option<&Value>) -> Option<(usize, Value)> { - messages? - .as_array()? - .iter() - .enumerate() - .rev() - .find(|(_, message)| message.get("role").and_then(Value::as_str) == Some("user")) - .and_then(|(index, message)| message.get("content").cloned().map(|value| (index, value))) -} - -#[allow(clippy::too_many_arguments)] -fn record_residual( - row: &Map, - relative_path: &str, - ordinal: usize, - step_id: i64, +fn request_message_context( + mut messages: Option, user_message_index: Option, - user_turn_id: Option, output_location: OutputLocation, - env_state: Option<&Value>, -) -> Value { - let mut residual = row.clone(); - for key in ["session_id", "step_id", "messages", "response"] { - residual.remove(key); - } - - let id_original = residual.remove("id"); - let id_present = id_original.is_some(); - let id_normalized = row - .get("id") - .and_then(Value::as_str) - .filter(|value| !value.is_empty()) - .map(str::to_string) - .unwrap_or_else(|| format!("step-{step_id}")); - let model_key = ["agent_model", "llm_model"] - .into_iter() - .find(|key| row.get(*key).and_then(Value::as_str).is_some()) - .map(str::to_string); - if let Some(key) = &model_key { - residual.remove(key); - } - let run_key = ["run_id", "run_bucket", "job_id"] - .into_iter() - .find(|key| { - row.get(*key) - .and_then(Value::as_str) - .is_some_and(|value| !value.is_empty()) - }) - .map(str::to_string); - if let Some(key) = &run_key { - residual.remove(key); - } - - let metric_fields = ROW_METRIC_FIELDS - .iter() - .filter(|field| row.contains_key(**field)) - .map(|field| Value::String((*field).to_string())) - .collect::>(); - for field in ROW_METRIC_FIELDS { - residual.remove(*field); - } - - let timestamp_from_env = env_state - .and_then(|value| value.get("created_at")) - .and_then(Value::as_str) - .is_some(); - let (created_at_kind, created_at_original, created_at_normalized) = if timestamp_from_env { - (None, None, None) - } else { - row.get("created_at").map_or((None, None, None), |value| { - residual.remove("created_at"); - let kind = match value { - Value::String(_) => "string", - Value::Number(number) if number.is_i64() || number.is_u64() => "integer", - Value::Number(_) => "float", - _ => "other", - }; - ( - Some(kind), - Some(value.clone()), - normalize_timestamp(value).map(Value::String), - ) - }) - }; - - let mut messages = row.get("messages").cloned(); - if let Some(values) = messages.as_mut().and_then(Value::as_array_mut) { - if let Some(index) = user_message_index { - if let Some(message) = values.get_mut(index).and_then(Value::as_object_mut) { - message.remove("content"); - } - } - if let OutputLocation::Message(index) = output_location { - if let Some(message) = values.get_mut(index).and_then(Value::as_object_mut) { - message.remove("content"); - if parse_tool_calls(message.get("tool_calls")).is_some() { - message.remove("tool_calls"); - } - } - } - } - let mut response = row.get("response").cloned(); - if matches!(output_location, OutputLocation::Response) { - if let Some(message) = response.as_mut().and_then(Value::as_object_mut) { +) -> Option { + let values = messages.as_mut()?.as_array_mut()?; + if let Some(index) = user_message_index { + if let Some(message) = values.get_mut(index).and_then(Value::as_object_mut) { message.remove("content"); - if parse_tool_calls(message.get("tool_calls")).is_some() { - message.remove("tool_calls"); - } - } - } - - let (output_kind, output_index) = match output_location { - OutputLocation::Response => ("response", None), - OutputLocation::Message(index) => ("message", Some(index)), - }; - json!({ - "relative_path": relative_path, - "ordinal": ordinal, - "step_id": step_id, - "user_message_index": user_message_index, - "user_turn_id": user_turn_id, - "output_kind": output_kind, - "output_index": output_index, - "id_present": id_present, - "id_original": id_original, - "id_normalized": id_normalized, - "model_key": model_key, - "run_key": run_key, - "metric_fields": metric_fields, - "created_at_kind": created_at_kind, - "created_at_original": created_at_original, - "created_at_normalized": created_at_normalized, - "messages": messages, - "response": response, - "residual": residual, - }) -} - -fn recover_record( - story: &StorylineDocument, - agent_turn: &StorylineTurn, - metadata: &Map, -) -> Result { - let mut record = metadata - .get("residual") - .and_then(Value::as_object) - .cloned() - .ok_or_else(|| anyhow::anyhow!("OpenAI record residual must be an object"))?; - let step_id = metadata - .get("step_id") - .and_then(Value::as_i64) - .ok_or_else(|| anyhow::anyhow!("OpenAI record residual missing step_id"))?; - insert_authoritative( - &mut record, - "session_id", - Value::String(story.session_id.clone()), - "record", - ); - insert_authoritative(&mut record, "step_id", json!(step_id), "record"); - - if metadata - .get("id_present") - .and_then(Value::as_bool) - .unwrap_or(false) - { - let call_id = agent_turn - .extra - .as_ref() - .and_then(|value| value.get("call_id")) - .and_then(Value::as_str); - let normalized = metadata.get("id_normalized").and_then(Value::as_str); - let value = if call_id == normalized { - metadata.get("id_original").cloned().unwrap_or(Value::Null) - } else { - call_id.map_or(Value::Null, |value| Value::String(value.to_string())) - }; - insert_authoritative(&mut record, "id", value, "record"); - } - if let Some(key) = metadata.get("model_key").and_then(Value::as_str) { - if let Some(model) = &agent_turn.model_name { - insert_authoritative(&mut record, key, Value::String(model.clone()), "record"); - } - } - if let Some(key) = metadata.get("run_key").and_then(Value::as_str) { - if let Some(run_id) = &story.run_id { - insert_authoritative(&mut record, key, Value::String(run_id.clone()), "record"); } } - if let Some(fields) = metadata.get("metric_fields").and_then(Value::as_array) { - for field in fields.iter().filter_map(Value::as_str) { - if let Some(value) = agent_turn - .metrics - .as_ref() - .and_then(|value| value.get(field)) - { - insert_authoritative(&mut record, field, value.clone(), "record"); - } + if let OutputLocation::Message(index) = output_location { + if index < values.len() { + values.remove(index); } } - if let Some(kind) = metadata.get("created_at_kind").and_then(Value::as_str) { - let encoded = match agent_turn.timestamp.as_deref() { - Some(timestamp) - if metadata - .get("created_at_normalized") - .and_then(Value::as_str) - == Some(timestamp) => - { - metadata - .get("created_at_original") - .cloned() - .unwrap_or(encode_timestamp(timestamp, kind)?) - } - Some(timestamp) => encode_timestamp(timestamp, kind)?, - None => metadata - .get("created_at_original") - .cloned() - .unwrap_or(Value::Null), - }; - insert_authoritative(&mut record, "created_at", encoded, "record"); - } - - let user_turn = metadata - .get("user_turn_id") - .and_then(Value::as_i64) - .and_then(|id| story.turns.iter().find(|turn| turn.id == id)); - let output_kind = metadata - .get("output_kind") - .and_then(Value::as_str) - .ok_or_else(|| anyhow::anyhow!("OpenAI record residual missing output_kind"))?; - let output_index = metadata.get("output_index").and_then(Value::as_u64); - - if let Some(mut messages) = metadata.get("messages").filter(|v| !v.is_null()).cloned() { - let values = messages - .as_array_mut() - .ok_or_else(|| anyhow::anyhow!("OpenAI messages residual must be an array"))?; - if let (Some(index), Some(user_turn)) = ( - metadata - .get("user_message_index") - .and_then(Value::as_u64) - .map(|value| value as usize), - user_turn, - ) { - let message = values - .get_mut(index) - .and_then(Value::as_object_mut) - .ok_or_else(|| anyhow::anyhow!("OpenAI user message residual is invalid"))?; - message.insert("content".into(), user_turn.message.clone()); - } - if output_kind == "message" { - let index = output_index - .ok_or_else(|| anyhow::anyhow!("OpenAI output message index is missing"))? - as usize; - let message = values - .get_mut(index) - .and_then(Value::as_object_mut) - .ok_or_else(|| anyhow::anyhow!("OpenAI output message residual is invalid"))?; - apply_output(message, agent_turn)?; - } - insert_authoritative(&mut record, "messages", messages, "record"); + values.retain(|message| message.get("role").and_then(Value::as_str) != Some("tool")); + for message in values.iter_mut().filter_map(Value::as_object_mut) { + retain_canonical_openai_message(message); } + messages +} - if let Some(mut response) = metadata.get("response").filter(|v| !v.is_null()).cloned() { - if output_kind == "response" { - let message = response - .as_object_mut() - .ok_or_else(|| anyhow::anyhow!("OpenAI response residual must be an object"))?; - apply_output(message, agent_turn)?; +fn retain_canonical_openai_message(message: &mut Map) { + message.retain(|key, _| { + matches!( + key.as_str(), + "role" | "content" | "name" | "tool_call_id" | "tool_calls" + ) + }); + let Some(calls) = message.get_mut("tool_calls").and_then(Value::as_array_mut) else { + return; + }; + for call in calls.iter_mut().filter_map(Value::as_object_mut) { + call.retain(|key, _| matches!(key.as_str(), "id" | "type" | "function")); + if let Some(function) = call.get_mut("function").and_then(Value::as_object_mut) { + function.retain(|key, _| matches!(key.as_str(), "name" | "arguments")); } - insert_authoritative(&mut record, "response", response, "record"); } - Ok(Value::Object(record)) } -fn apply_output(message: &mut Map, turn: &StorylineTurn) -> Result<()> { - message.insert("content".into(), turn.message.clone()); - if let Some(calls) = &turn.tool_calls { - message.insert("tool_calls".into(), encode_tool_calls(calls)?); - } - Ok(()) +fn last_user_message(messages: Option<&Value>) -> Option<(usize, Value)> { + messages? + .as_array()? + .iter() + .enumerate() + .rev() + .find(|(_, message)| message.get("role").and_then(Value::as_str) == Some("user")) + .and_then(|(index, message)| message.get("content").cloned().map(|value| (index, value))) } -fn insert_authoritative(target: &mut Map, key: &str, value: Value, scope: &str) { - if target.contains_key(key) { - tracing::warn!( - source_format = "openai-msg", - source_key = key, - target_key = key, - scope, - "OpenAI residual conflicts with an authoritative Storyline field" - ); - } - target.insert(key.to_string(), value); +fn parse_tool_results(messages: Option<&Value>) -> Option { + let results = messages? + .as_array()? + .iter() + .filter(|message| message.get("role").and_then(Value::as_str) == Some("tool")) + .filter_map(|message| { + let source_call_id = message.get("tool_call_id")?.as_str()?; + if source_call_id.is_empty() { + return None; + } + Some(json!({ + "source_call_id": source_call_id, + "content": message.get("content").cloned().unwrap_or(Value::Null), + })) + }) + .collect::>(); + (!results.is_empty()).then(|| json!({"results": results})) } -fn encode_timestamp(timestamp: &str, kind: &str) -> Result { - if kind == "string" { - return Ok(Value::String(timestamp.to_string())); - } - if kind == "other" { - return Ok(Value::String(timestamp.to_string())); - } - let parsed = chrono::DateTime::parse_from_rfc3339(timestamp)?; - if kind == "integer" && parsed.timestamp_subsec_nanos() == 0 { - Ok(json!(parsed.timestamp())) - } else { - Ok(json!( - parsed.timestamp() as f64 - + f64::from(parsed.timestamp_subsec_nanos()) / 1_000_000_000.0 - )) - } +fn encode_tool_results(observation: Option<&Value>) -> Vec { + observation + .and_then(|value| value.get("results")) + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(|result| { + let source_call_id = result.get("source_call_id")?.as_str()?; + if source_call_id.is_empty() { + return None; + } + Some(json!({ + "role": "tool", + "tool_call_id": source_call_id, + "content": result.get("content").cloned().unwrap_or(Value::Null), + })) + }) + .collect() } fn validate_relative_path(path: &Path) -> Result { @@ -964,28 +990,13 @@ fn parse_tool_calls(value: Option<&Value>) -> Option> { } _ => arguments, }; - let mut call_residual = call.clone(); - call_residual.remove("id"); - call_residual.remove("type"); - call_residual.remove("function"); - let mut function_residual = function.clone(); - function_residual.remove("name"); - let raw_arguments = function_residual.remove("arguments"); Some(StorylineToolCall { tool_call_id, function_name, arguments, result: Default::default(), duration_ms: None, - extra: Some(json!({ - OPENAI_EXTENSION_KEY: { - "kind": "tool_call", - "type": call.get("type"), - "call": call_residual, - "function": function_residual, - "arguments_were_string": raw_arguments.is_some_and(|value| value.is_string()), - } - })), + extra: None, }) }) .collect::>(); @@ -996,38 +1007,14 @@ fn encode_tool_calls(calls: &[StorylineToolCall]) -> Result { calls .iter() .map(|call| { - let metadata = call - .extra - .as_ref() - .and_then(|value| value.get(OPENAI_EXTENSION_KEY)) - .and_then(Value::as_object); - let mut output = metadata - .and_then(|value| value.get("call")) - .and_then(Value::as_object) - .cloned() - .unwrap_or_default(); - output.insert("id".into(), Value::String(call.tool_call_id.clone())); - if let Some(kind) = metadata.and_then(|value| value.get("type")) { - output.insert("type".into(), kind.clone()); - } - let mut function = metadata - .and_then(|value| value.get("function")) - .and_then(Value::as_object) - .cloned() - .unwrap_or_default(); - function.insert("name".into(), Value::String(call.function_name.clone())); - let arguments = if metadata - .and_then(|value| value.get("arguments_were_string")) - .and_then(Value::as_bool) - .unwrap_or(false) - { - Value::String(serde_json::to_string(&call.arguments)?) - } else { - call.arguments.clone() - }; - function.insert("arguments".into(), arguments); - output.insert("function".into(), Value::Object(function)); - Ok(Value::Object(output)) + Ok(json!({ + "id": call.tool_call_id, + "type": "function", + "function": { + "name": call.function_name, + "arguments": serde_json::to_string(&call.arguments)?, + } + })) }) .collect::>>() .map(Value::Array) @@ -1150,6 +1137,30 @@ fn number_to_i64(value: &Value) -> Option { #[cfg(test)] mod tests { use super::*; + + #[test] + fn openai_unknown_fields_use_exact_row_paths() { + let input = json!({"root_vendor": 1, "session_steps": [{ + "session_id": "s", "step_id": 1, + "messages": [{ + "role": "user", "content": "hi", "message_vendor": null, "0": true + }], + "response": {"role": "assistant", "content": "ok"}, + "row_vendor": [3, 2, 1] + }]}); + let stories = parse_openai_msg_corpus_value(&input, "corpus.json").unwrap(); + let fields = &stories[0].unknown_fields.sources["openai-msg"].fields; + assert_eq!(fields["/root_vendor"], 1); + assert_eq!(fields["/session_steps/0/row_vendor"], json!([3, 2, 1])); + assert_eq!( + fields["/session_steps/0/messages/0/message_vendor"], + Value::Null + ); + assert_eq!( + stories[0].unknown_key_counts["openai-msg"]["/session_steps/*/messages/*/0"], + 1 + ); + } #[cfg(feature = "lance-store")] use crate::store::StorylineLanceStore; @@ -1201,7 +1212,7 @@ mod tests { } #[test] - fn corpus_roundtrip_is_json_semantically_lossless() { + fn corpus_roundtrip_emits_canonical_envelope_in_source_order() { let input = corpus(); let stories = parse_openai_msg_corpus_value(&input, "corpus.json").unwrap(); assert_eq!(stories.len(), 2); @@ -1217,7 +1228,12 @@ mod tests { let recovered = recover_openai_msg_files(&stories).unwrap(); assert_eq!(recovered.len(), 1); assert_eq!(recovered[0].relative_path, PathBuf::from("corpus.json")); - assert_eq!(recovered[0].document, input); + let rows = recovered[0].document["session_steps"].as_array().unwrap(); + assert_eq!(rows.len(), 3); + assert_eq!(rows[0]["session_id"], "s-1"); + assert_eq!(rows[0]["step_id"], 2); + assert_eq!(rows[1]["session_id"], "s-2"); + assert_eq!(rows[2]["step_id"], 1); } #[test] @@ -1238,7 +1254,7 @@ mod tests { } #[test] - fn openai_residual_preserves_unknowns_but_storyline_content_is_authoritative() { + fn openai_unknown_fields_preserve_values_but_storyline_content_is_authoritative() { let input = corpus(); let mut stories = parse_openai_msg_corpus_value(&input, "corpus.json").unwrap(); assert!(!serde_json::to_string(&stories) @@ -1248,14 +1264,55 @@ mod tests { stories[0].turns[0].message = json!("edited user"); stories[0].turns[1].message = json!("edited assistant"); let recovered = recover_openai_msg_files(&stories).unwrap(); - let rows = recovered[0].document.as_array().unwrap(); + let rows = recovered[0].document["session_steps"].as_array().unwrap(); let first_session_row = rows.iter().find(|row| row["id"] == "evt-1").unwrap(); assert_eq!(first_session_row["messages"][1]["content"], "edited user"); + assert_eq!(first_session_row["response"]["content"], "edited assistant"); + assert_eq!(rows[0]["unknown"], Value::Null); + } + + #[test] + fn message_unknowns_and_tool_results_restore_once() { + let input = json!({"session_steps": [{ + "id": "", + "session_id": "s", + "agent_id": "agent", + "step_id": 1, + "messages": [ + {"role": "user", "content": "run", "vendor_message": 7}, + {"role": "tool", "tool_call_id": "call-1", "content": "ok"} + ], + "response": { + "role": "assistant", + "content": "done", + "tool_calls": [{ + "id": "call-1", + "type": "function", + "function": {"name": "inspect", "arguments": "{}"} + }] + } + }]}); + + let stories = parse_openai_msg_corpus_value(&input, "tool-results.json").unwrap(); assert_eq!( - first_session_row["messages"][2]["content"], - "edited assistant" + stories[0].turns[1].observation.as_ref().unwrap()["results"][0], + json!({"source_call_id": "call-1", "content": "ok"}) ); - assert_eq!(rows[0]["unknown"], Value::Null); + let recovered = recover_openai_msg_files(&stories).unwrap(); + let row = &recovered[0].document["session_steps"][0]; + assert_eq!(row["agent_id"], "agent"); + assert_eq!(row["id"], ""); + assert_eq!(row["messages"][0]["vendor_message"], 7); + assert_eq!( + row["messages"] + .as_array() + .unwrap() + .iter() + .filter(|message| message["role"] == "tool") + .count(), + 1 + ); + assert_ne!(row["created_at"], ""); } #[test] @@ -1267,7 +1324,9 @@ mod tests { }); let stories = parse_openai_msg_corpus_value(&input, "session_steps.json").unwrap(); let recovered = recover_openai_msg_files(&stories).unwrap(); - assert_eq!(recovered[0].document, input); + assert_eq!(recovered[0].document["custom"], Value::Null); + assert_eq!(recovered[0].document["session_id"], "s-1"); + assert!(recovered[0].document["session_steps"].is_array()); } #[test] @@ -1286,10 +1345,12 @@ mod tests { assert_eq!(stories[0].turns.len(), 2); assert_eq!(stories[0].turns[0].source, "user"); assert_eq!(stories[0].turns[1].source, "agent"); - assert_eq!( - recover_openai_msg_files(&stories).unwrap()[0].document, - input - ); + let recovered = recover_openai_msg_files(&stories).unwrap(); + let row = &recovered[0].document["session_steps"][0]; + assert_eq!(row["session_id"], "child-session"); + assert_eq!(row["step_id"], 7); + assert_eq!(row["messages"][0]["content"], "question"); + assert_eq!(row["response"]["content"], "answer"); } #[test] @@ -1338,6 +1399,9 @@ mod tests { async fn corpus_import_and_recovery_roundtrip_through_lance() { let input = corpus(); let expected = parse_openai_msg_corpus_value(&input, "corpus.json").unwrap(); + let canonical = recover_openai_msg_files(&expected).unwrap()[0] + .document + .clone(); let temporary = tempfile::tempdir().unwrap(); let store = StorylineLanceStore::open(temporary.path()).await.unwrap(); store.replace_storylines(&expected).await.unwrap(); @@ -1357,7 +1421,7 @@ mod tests { assert_eq!(recovered.len(), 1); assert_eq!(recovered[0].relative_path, PathBuf::from("corpus.json")); - assert_eq!(recovered[0].document, input); + assert_eq!(recovered[0].document, canonical); } #[cfg(feature = "lance-store")] @@ -1366,6 +1430,9 @@ mod tests { let mut input = corpus(); input[0]["created_at"] = json!(1_700_000_001.123_456_f64); let expected = parse_openai_msg_corpus_value(&input, "fractional.json").unwrap(); + let canonical = recover_openai_msg_files(&expected).unwrap()[0] + .document + .clone(); let temporary = tempfile::tempdir().unwrap(); let store = StorylineLanceStore::open(temporary.path()).await.unwrap(); store.replace_storylines(&expected).await.unwrap(); @@ -1384,7 +1451,7 @@ mod tests { assert_eq!( recover_openai_msg_files(&restored).unwrap()[0].document, - input + canonical ); } @@ -1399,6 +1466,9 @@ mod tests { "messages": [{"role": "assistant", "content": "ok"}] }]); let expected = parse_openai_msg_corpus_value(&input, "nulls.json").unwrap(); + let canonical = recover_openai_msg_files(&expected).unwrap()[0] + .document + .clone(); let temporary = tempfile::tempdir().unwrap(); let store = StorylineLanceStore::open(temporary.path()).await.unwrap(); store.replace_storylines(&expected).await.unwrap(); @@ -1417,7 +1487,7 @@ mod tests { assert_eq!( recover_openai_msg_files(&restored).unwrap()[0].document, - input + canonical ); } } diff --git a/crates/persisting-pchronicle/src/formats/storyline.rs b/crates/persisting-pchronicle/src/formats/storyline.rs index b30e1a41..a684cfcc 100644 --- a/crates/persisting-pchronicle/src/formats/storyline.rs +++ b/crates/persisting-pchronicle/src/formats/storyline.rs @@ -4,158 +4,14 @@ //! Short wire keys (`src`, `msg`, `ts`, …); timing convenience fields //! (`latency_ms` / `ttft_ms` / `duration_ms`) lift common metrics. -use std::collections::{BTreeMap, BTreeSet}; +use std::collections::BTreeMap; use serde::{Deserialize, Serialize}; use serde_json::Value; +use super::unknown_fields::{compute_unknown_key_counts, StorylineUnknownFields, UnknownKeyCounts}; use crate::{InputIssue, InputResult, Result}; -/// Presence semantics for interchange fields where missing and explicit null -/// carry different meanings. -#[derive(Debug, Clone, Default, PartialEq)] -pub enum FieldPresence { - #[default] - Missing, - Null, - Value(T), -} - -impl FieldPresence { - pub fn is_missing(&self) -> bool { - matches!(self, Self::Missing) - } - - pub fn is_null(&self) -> bool { - matches!(self, Self::Null) - } - - pub fn value(&self) -> Option<&T> { - match self { - Self::Value(value) => Some(value), - Self::Missing | Self::Null => None, - } - } - - pub fn as_ref(&self) -> Option<&T> { - self.value() - } - - pub fn into_option(self) -> Option { - match self { - Self::Value(value) => Some(value), - Self::Missing | Self::Null => None, - } - } -} - -impl Serialize for FieldPresence { - fn serialize(&self, serializer: S) -> std::result::Result - where - S: serde::Serializer, - { - match self { - Self::Missing | Self::Null => serializer.serialize_none(), - Self::Value(value) => value.serialize(serializer), - } - } -} - -impl<'de, T: Deserialize<'de>> Deserialize<'de> for FieldPresence { - fn deserialize(deserializer: D) -> std::result::Result - where - D: serde::Deserializer<'de>, - { - Ok(match Option::::deserialize(deserializer)? { - Some(value) => Self::Value(value), - None => Self::Null, - }) - } -} - -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PresenceState { - Missing, - Null, - #[default] - Value, -} - -/// Shape of the physical document collection that contained a Storyline. -/// -/// This is format-neutral collection semantics, not an editable -/// format-specific residual. It allows collection shape and ordering to pass -/// through the authoritative Storyline model and Lance storage. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum StorylineCollectionShape { - Single, - Sequence, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum StorylineRootField { - TrajectoryId, - Notes, - FinalMetrics, - ContinuedTrajectoryRef, - Extra, - SubagentTrajectories, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum StorylineAgentField { - ModelName, - ToolDefinitions, - Extra, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum StorylineTurnField { - Timestamp, - ModelName, - ReasoningEffort, - ReasoningContent, - ToolCalls, - Observation, - Metrics, - Extra, - LlmCallCount, - IsCopiedContext, -} - -#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] -pub struct StorylinePresence { - #[serde(default, skip_serializing_if = "is_value_presence")] - pub session_id: PresenceState, - #[serde(default, skip_serializing_if = "BTreeSet::is_empty")] - pub root_nulls: BTreeSet, - #[serde(default, skip_serializing_if = "BTreeSet::is_empty")] - pub agent_nulls: BTreeSet, - #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] - pub turn_nulls: BTreeMap>, - #[serde(default, skip_serializing_if = "BTreeSet::is_empty")] - pub tool_call_extra_nulls: BTreeSet, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub collection_shape: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub collection_ordinal: Option, -} - -fn is_value_presence(value: &PresenceState) -> bool { - *value == PresenceState::Value -} - -impl StorylinePresence { - pub fn is_default(&self) -> bool { - self == &Self::default() - } -} - #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct StorylineDocument { #[serde(default, skip_serializing_if = "Option::is_none")] @@ -186,8 +42,10 @@ pub struct StorylineDocument { pub continued_trajectory_ref: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub extra: Option, - #[serde(default, skip_serializing_if = "StorylinePresence::is_default")] - pub presence: StorylinePresence, + #[serde(default, skip_serializing_if = "StorylineUnknownFields::is_empty")] + pub unknown_fields: StorylineUnknownFields, + #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] + pub unknown_key_counts: UnknownKeyCounts, pub turns: Vec, } @@ -298,8 +156,8 @@ pub struct StorylineToolCall { pub function_name: String, #[serde(rename = "args")] pub arguments: Value, - #[serde(default, skip_serializing_if = "FieldPresence::is_missing")] - pub result: FieldPresence, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub result: Option, /// Tool execution wall time in milliseconds. #[serde(default, skip_serializing_if = "Option::is_none")] pub duration_ms: Option, @@ -330,7 +188,8 @@ impl StorylineDocument { final_metrics: None, continued_trajectory_ref: None, extra: None, - presence: StorylinePresence::default(), + unknown_fields: StorylineUnknownFields::default(), + unknown_key_counts: UnknownKeyCounts::default(), turns: Vec::new(), } } @@ -361,23 +220,9 @@ impl StorylineDocument { if self.agent.id.is_empty() { return Err(InputIssue::invalid("storyline.agent.id is required")); } - if self - .presence - .collection_ordinal - .is_some_and(|ordinal| ordinal < 0) - { + if compute_unknown_key_counts(&self.unknown_fields)? != self.unknown_key_counts { return Err(InputIssue::invalid( - "storyline collection ordinal cannot be negative", - )); - } - if self.presence.collection_shape == Some(StorylineCollectionShape::Single) - && self - .presence - .collection_ordinal - .is_some_and(|ordinal| ordinal != 0) - { - return Err(InputIssue::invalid( - "single-document Storyline collection ordinal must be zero", + "storyline unknown_key_counts do not match unknown_fields", )); } let mut seen = std::collections::HashSet::new(); @@ -397,6 +242,11 @@ impl StorylineDocument { } Ok(()) } + + pub fn refresh_unknown_key_counts(&mut self) -> InputResult<()> { + self.unknown_key_counts = compute_unknown_key_counts(&self.unknown_fields)?; + Ok(()) + } } #[cfg(all(test, feature = "lance-store"))] @@ -408,22 +258,135 @@ pub fn parse_storyline_document(input: &str) -> Result { mod tests { use super::*; + fn story_with_source_normalized_counts() -> StorylineDocument { + let mut story = StorylineDocument::new("session", "agent"); + for (source, source_id, pointer) in [ + ("atif", "atif-doc", "/steps/0/vendor"), + ("atif", "atif-doc", "/steps/1/vendor"), + ("actf", "actf-doc", "/attempts/7/trajectory/steps/0/vendor"), + ("actf", "actf-doc", "/attempts/7/trajectory/steps/1/vendor"), + ( + "openai-msg", + "openai-doc", + "/session_steps/0/messages/0/vendor", + ), + ( + "openai-msg", + "openai-doc", + "/session_steps/1/messages/1/vendor", + ), + ("agenticmd", "agenticmd-doc", "/blocks/0/header/vendor"), + ("agenticmd", "agenticmd-doc", "/blocks/1/header/vendor"), + ("future-format", "future-doc", "/items/0/vendor"), + ("future-format", "future-doc", "/items/1/vendor"), + ] { + story + .unknown_fields + .insert(source, source_id, pointer, serde_json::json!(true)) + .unwrap(); + } + story.unknown_key_counts = BTreeMap::from([ + ( + "atif".into(), + BTreeMap::from([("/steps/*/vendor".into(), 2)]), + ), + ( + "actf".into(), + BTreeMap::from([("/attempts/7/trajectory/steps/*/vendor".into(), 2)]), + ), + ( + "openai-msg".into(), + BTreeMap::from([("/session_steps/*/messages/*/vendor".into(), 2)]), + ), + ( + "agenticmd".into(), + BTreeMap::from([("/blocks/*/header/vendor".into(), 2)]), + ), + ( + "future-format".into(), + BTreeMap::from([("/items/0/vendor".into(), 1), ("/items/1/vendor".into(), 1)]), + ), + ]); + story + } + + #[test] + fn validate_accepts_source_normalized_unknown_key_counts() { + story_with_source_normalized_counts().validate().unwrap(); + } + + #[test] + fn validate_rejects_stale_source_normalized_unknown_key_counts() { + let mut story = story_with_source_normalized_counts(); + *story + .unknown_key_counts + .get_mut("atif") + .unwrap() + .get_mut("/steps/*/vendor") + .unwrap() = 1; + assert!(story.validate().is_err()); + } + #[test] - fn legacy_long_field_names_are_rejected() { - let legacy = serde_json::json!({ + fn refresh_counts_wildcards_only_schema_array_positions() { + let mut story = StorylineDocument::new("session", "agent"); + for (source, source_id, pointer) in [ + ("atif", "atif-doc", "/steps/0/0"), + ("actf", "actf-doc", "/attempts/1/trajectory/steps/0/0"), + ("openai-msg", "openai-doc", "/session_steps/0/messages/0/0"), + ("agenticmd", "agenticmd-doc", "/frontmatter/0"), + ("agenticmd", "agenticmd-doc", "/blocks/0/header/0"), + ("future-format", "future-doc", "/items/0/0"), + ] { + story + .unknown_fields + .insert(source, source_id, pointer, serde_json::json!(true)) + .unwrap(); + } + + story.refresh_unknown_key_counts().unwrap(); + + assert_eq!(story.unknown_key_counts["atif"]["/steps/*/0"], 1); + assert_eq!( + story.unknown_key_counts["actf"]["/attempts/1/trajectory/steps/*/0"], + 1 + ); + assert_eq!( + story.unknown_key_counts["openai-msg"]["/session_steps/*/messages/*/0"], + 1 + ); + assert_eq!(story.unknown_key_counts["agenticmd"]["/frontmatter/0"], 1); + assert_eq!( + story.unknown_key_counts["agenticmd"]["/blocks/*/header/0"], + 1 + ); + assert_eq!(story.unknown_key_counts["future-format"]["/items/0/0"], 1); + story.validate().unwrap(); + } + + #[test] + fn storyline_serialization_omits_empty_unknown_fields() { + let story = StorylineDocument::new("session", "agent"); + let value = serde_json::to_value(story).unwrap(); + assert!(value.get("unknown_fields").is_none()); + } + + #[test] + fn unsupported_long_field_names_are_rejected() { + let document = serde_json::json!({ "session_id": "session-1", "agent": { "id": "agent-1" }, "turns": [] }); - assert!(serde_json::from_value::(legacy).is_err()); + assert!(serde_json::from_value::(document).is_err()); } #[test] - fn tool_result_presence_distinguishes_missing_null_and_value() { + fn tool_result_canonicalizes_missing_and_null() { let base = serde_json::json!({"tcid":"call-1","fn":"lookup","args":{}}); let missing: StorylineToolCall = serde_json::from_value(base.clone()).unwrap(); - assert_eq!(missing.result, FieldPresence::Missing); + assert_eq!(missing.result, None); assert!(serde_json::to_value(missing) .unwrap() .get("result") @@ -432,16 +395,13 @@ mod tests { let mut null = base.clone(); null["result"] = Value::Null; let null: StorylineToolCall = serde_json::from_value(null).unwrap(); - assert_eq!(null.result, FieldPresence::Null); - assert_eq!(serde_json::to_value(null).unwrap()["result"], Value::Null); + assert_eq!(null.result, None); + assert!(serde_json::to_value(null).unwrap().get("result").is_none()); let mut value = base; value["result"] = serde_json::json!({"answer": 42}); let value: StorylineToolCall = serde_json::from_value(value).unwrap(); - assert_eq!( - value.result, - FieldPresence::Value(serde_json::json!({"answer": 42})) - ); + assert_eq!(value.result, Some(serde_json::json!({"answer": 42}))); assert_eq!( serde_json::to_value(value).unwrap()["result"], serde_json::json!({"answer": 42}) diff --git a/crates/persisting-pchronicle/src/formats/unknown_fields.rs b/crates/persisting-pchronicle/src/formats/unknown_fields.rs new file mode 100644 index 00000000..805516b4 --- /dev/null +++ b/crates/persisting-pchronicle/src/formats/unknown_fields.rs @@ -0,0 +1,1081 @@ +use crate::format::DocumentFormat; +use crate::formats::StorylineDocument; +use crate::{InputIssue, InputResult, Result}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::{btree_map::Entry, BTreeMap, BTreeSet}; + +pub const DEFAULT_MAX_UNKNOWN_FIELDS: usize = 4096; +pub const DEFAULT_MAX_UNKNOWN_BYTES: usize = 1024 * 1024; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct UnknownFieldLimits { + pub max_fields: usize, + pub max_bytes: usize, +} + +impl Default for UnknownFieldLimits { + fn default() -> Self { + Self { + max_fields: DEFAULT_MAX_UNKNOWN_FIELDS, + max_bytes: DEFAULT_MAX_UNKNOWN_BYTES, + } + } +} + +impl UnknownFieldLimits { + pub fn validate(self) -> InputResult<()> { + if self.max_fields == 0 || self.max_fields == usize::MAX { + return Err(InputIssue::invalid( + "unknown field limit must be a finite positive value", + )); + } + if self.max_bytes == 0 || self.max_bytes == usize::MAX { + return Err(InputIssue::invalid( + "unknown byte limit must be a finite positive value", + )); + } + Ok(()) + } +} + +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +pub struct SourceUnknownFields { + pub source_document_id: String, + pub fields: BTreeMap, +} + +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +pub struct StorylineUnknownFields { + pub sources: BTreeMap, +} + +pub type UnknownFieldCounts = BTreeMap; +pub type UnknownKeyCounts = BTreeMap; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct CarrierBinding { + pub story_index: usize, + pub pointer: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +struct StorylineEnvelopeWire { + unknown_fields: UnknownFieldsEnvelopeWire, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +struct UnknownFieldsEnvelopeWire { + version: u32, + by_trajectory: BTreeMap, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +struct UnknownFieldsCarrierWire { + sources: BTreeMap, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +struct SourceUnknownFieldsWire { + source_document_id: String, + fields: BTreeMap, +} + +pub(crate) fn take_unknown_fields_envelope( + document: &mut Value, +) -> InputResult> { + let Some(raw_envelope) = document + .as_object() + .and_then(|object| object.get("_storyline")) + .cloned() + else { + return Ok(BTreeMap::new()); + }; + + let envelope: StorylineEnvelopeWire = serde_json::from_value(raw_envelope) + .map_err(|error| InputIssue::invalid(format!("invalid _storyline envelope: {error}")))?; + if envelope.unknown_fields.version != 1 { + return Err(InputIssue::invalid(format!( + "unsupported _storyline unknown_fields version {}; expected 1", + envelope.unknown_fields.version + ))); + } + + let mut carried = BTreeMap::new(); + for (carrier, fields) in envelope.unknown_fields.by_trajectory { + validate_json_pointer(&carrier)?; + let fields = StorylineUnknownFields { + sources: fields + .sources + .into_iter() + .map(|(source, fields)| { + ( + source, + SourceUnknownFields { + source_document_id: fields.source_document_id, + fields: fields.fields, + }, + ) + }) + .collect(), + }; + compute_unknown_key_counts(&fields)?; + carried.insert(carrier, fields); + } + + document + .as_object_mut() + .expect("an object containing _storyline remains an object") + .remove("_storyline"); + Ok(carried) +} + +pub(crate) fn attach_carried_unknown_fields( + envelope: BTreeMap, + carriers: &[CarrierBinding], + stories: &mut [StorylineDocument], + limits: UnknownFieldLimits, +) -> InputResult<()> { + limits.validate()?; + let by_pointer = validate_carrier_bindings(carriers, stories.len())?; + + let mut merged = stories + .iter() + .map(|story| story.unknown_fields.clone()) + .collect::>(); + for (pointer, carried) in envelope { + validate_json_pointer(&pointer)?; + let story_index = by_pointer.get(&pointer).ok_or_else(|| { + InputIssue::invalid(format!( + "_storyline envelope carrier '{pointer}' is not bound to a trajectory" + )) + })?; + merge_unknown_fields(&mut merged[*story_index], carried)?; + } + + let counts = merged + .iter() + .map(|fields| validate_unknown_fields(fields, limits)) + .collect::>>()?; + for ((story, fields), counts) in stories.iter_mut().zip(merged).zip(counts) { + story.unknown_fields = fields; + story.unknown_key_counts = counts; + } + Ok(()) +} + +pub(crate) fn write_foreign_unknown_fields_envelope( + target_format: DocumentFormat, + document: &mut Value, + stories: &[StorylineDocument], + carriers: &[CarrierBinding], +) -> Result<()> { + if document + .as_object() + .is_some_and(|object| object.contains_key("_storyline")) + { + anyhow::bail!("target document already contains reserved key '_storyline'"); + } + + let by_pointer = validate_carrier_bindings(carriers, stories.len())?; + for pointer in by_pointer.keys() { + if document.pointer(pointer).is_none() { + anyhow::bail!("carrier JSON Pointer '{pointer}' does not exist in target document"); + } + } + + let bound_story_indexes = carriers + .iter() + .map(|carrier| carrier.story_index) + .collect::>(); + let target_source = target_format.as_str(); + for (story_index, story) in stories.iter().enumerate() { + let has_foreign_fields = story + .unknown_fields + .sources + .iter() + .any(|(source, fields)| source != target_source && !fields.fields.is_empty()); + if has_foreign_fields && !bound_story_indexes.contains(&story_index) { + anyhow::bail!( + "Storyline at index {story_index} has foreign unknown fields but no carrier binding" + ); + } + } + + let mut by_trajectory = BTreeMap::new(); + for carrier in carriers { + let story = &stories[carrier.story_index]; + let sources = story + .unknown_fields + .sources + .iter() + .filter(|(source, fields)| { + source.as_str() != target_source && !fields.fields.is_empty() + }) + .map(|(source, fields)| { + for pointer in fields.fields.keys() { + validate_json_pointer(pointer)?; + } + Ok(( + source.clone(), + SourceUnknownFieldsWire { + source_document_id: fields.source_document_id.clone(), + fields: fields.fields.clone(), + }, + )) + }) + .collect::>>()?; + if !sources.is_empty() { + by_trajectory.insert( + carrier.pointer.clone(), + UnknownFieldsCarrierWire { sources }, + ); + } + } + + if by_trajectory.is_empty() { + return Ok(()); + } + let object = document + .as_object_mut() + .ok_or_else(|| anyhow::anyhow!("foreign unknown fields require an object document root"))?; + let envelope = StorylineEnvelopeWire { + unknown_fields: UnknownFieldsEnvelopeWire { + version: 1, + by_trajectory, + }, + }; + object.insert("_storyline".into(), serde_json::to_value(envelope)?); + Ok(()) +} + +fn validate_carrier_bindings( + carriers: &[CarrierBinding], + story_count: usize, +) -> InputResult> { + let mut by_pointer = BTreeMap::new(); + for carrier in carriers { + validate_json_pointer(&carrier.pointer)?; + if carrier.story_index >= story_count { + return Err(InputIssue::invalid(format!( + "carrier '{}' references missing Storyline index {}", + carrier.pointer, carrier.story_index + ))); + } + if by_pointer + .insert(carrier.pointer.clone(), carrier.story_index) + .is_some() + { + return Err(InputIssue::invalid(format!( + "duplicate carrier binding '{}'", + carrier.pointer + ))); + } + } + Ok(by_pointer) +} + +fn merge_unknown_fields( + target: &mut StorylineUnknownFields, + incoming: StorylineUnknownFields, +) -> InputResult<()> { + for (source, incoming_source) in incoming.sources { + match target.sources.entry(source.clone()) { + Entry::Vacant(entry) => { + entry.insert(incoming_source); + } + Entry::Occupied(mut entry) => { + let existing = entry.get_mut(); + if existing.source_document_id != incoming_source.source_document_id { + return Err(InputIssue::invalid(format!( + "unknown fields source '{source}' cannot change source_document_id" + ))); + } + for (pointer, value) in incoming_source.fields { + match existing.fields.entry(pointer.clone()) { + Entry::Vacant(entry) => { + entry.insert(value); + } + Entry::Occupied(entry) if entry.get() == &value => {} + Entry::Occupied(_) => { + return Err(InputIssue::invalid(format!( + "unknown fields source '{source}' has conflicting values at '{pointer}'" + ))) + } + } + } + } + } + } + Ok(()) +} + +impl StorylineUnknownFields { + pub fn is_empty(&self) -> bool { + self.sources.values().all(|source| source.fields.is_empty()) + } + + pub fn insert( + &mut self, + source: impl Into, + source_document_id: impl Into, + pointer: impl Into, + value: Value, + ) -> InputResult<()> { + let source = source.into(); + let source_document_id = source_document_id.into(); + let pointer = pointer.into(); + validate_json_pointer(&pointer)?; + + match self.sources.get_mut(&source) { + Some(existing) => { + if existing.source_document_id != source_document_id { + return Err(InputIssue::invalid(format!( + "unknown fields source '{source}' cannot change source_document_id" + ))); + } + existing.fields.insert(pointer, value); + } + None => { + self.sources.insert( + source, + SourceUnknownFields { + source_document_id, + fields: BTreeMap::from([(pointer, value)]), + }, + ); + } + } + Ok(()) + } + + pub fn validate_with( + &self, + limits: UnknownFieldLimits, + normalize: F, + ) -> InputResult + where + F: FnMut(&str, &str) -> InputResult, + { + validate_unknown_fields_with(self, limits, normalize) + } +} + +pub fn validate_json_pointer(pointer: &str) -> InputResult<()> { + decode_json_pointer(pointer).map(|_| ()) +} + +pub fn compute_unknown_key_counts( + fields: &StorylineUnknownFields, +) -> InputResult { + compute_unknown_key_counts_with(fields, normalize_unknown_pointer) +} + +pub fn validate_unknown_fields( + fields: &StorylineUnknownFields, + limits: UnknownFieldLimits, +) -> InputResult { + validate_unknown_fields_with(fields, limits, normalize_unknown_pointer) +} + +pub fn validate_unknown_fields_with( + fields: &StorylineUnknownFields, + limits: UnknownFieldLimits, + normalize: F, +) -> InputResult +where + F: FnMut(&str, &str) -> InputResult, +{ + limits.validate()?; + + let (field_count, byte_count) = logical_size(fields)?; + if field_count > limits.max_fields { + return Err(InputIssue::invalid(format!( + "unknown field count {field_count} exceeds configured limit {}", + limits.max_fields + ))); + } + if byte_count > limits.max_bytes { + return Err(InputIssue::invalid(format!( + "unknown field byte size {byte_count} exceeds configured limit {}", + limits.max_bytes + ))); + } + + compute_unknown_key_counts_with(fields, normalize) +} + +fn logical_size(fields: &StorylineUnknownFields) -> InputResult<(usize, usize)> { + let mut field_count = 0usize; + let mut byte_count = 0usize; + for source in fields.sources.values() { + byte_count = checked_size_add(byte_count, source.source_document_id.len())?; + for (pointer, value) in &source.fields { + field_count = field_count + .checked_add(1) + .ok_or_else(|| InputIssue::invalid("unknown field count overflow"))?; + byte_count = checked_size_add(byte_count, pointer.len())?; + byte_count = checked_size_add( + byte_count, + serde_json::to_vec(value) + .map_err(|error| InputIssue::invalid(error.to_string()))? + .len(), + )?; + } + } + Ok((field_count, byte_count)) +} + +fn checked_size_add(total: usize, additional: usize) -> InputResult { + total + .checked_add(additional) + .ok_or_else(|| InputIssue::invalid("unknown field byte count overflow")) +} + +fn compute_unknown_key_counts_with( + fields: &StorylineUnknownFields, + mut normalize: F, +) -> InputResult +where + F: FnMut(&str, &str) -> InputResult, +{ + let mut counts = UnknownKeyCounts::new(); + for (source, source_fields) in &fields.sources { + let source_counts = counts.entry(source.clone()).or_default(); + for pointer in source_fields.fields.keys() { + validate_json_pointer(pointer)?; + let normalized_pointer = normalize(source, pointer)?; + let count = source_counts.entry(normalized_pointer).or_default(); + *count = count.saturating_add(1); + } + } + Ok(counts) +} + +fn normalize_unknown_pointer(source: &str, pointer: &str) -> InputResult { + match source { + "atif" => normalize_atif_pointer(source, pointer), + "actf" => normalize_actf_pointer(source, pointer), + "openai-msg" => normalize_openai_pointer(source, pointer), + "agenticmd" => normalize_agenticmd_unknown_pointer(source, pointer), + _ => { + validate_json_pointer(pointer)?; + Ok(pointer.to_owned()) + } + } +} + +pub(crate) fn normalize_atif_pointer(source: &str, pointer: &str) -> InputResult { + let mut tokens = decode_json_pointer(pointer)?; + if source != "atif" { + return Ok(encode_json_pointer(&tokens)); + } + + fn normalize_trajectory(tokens: &mut [String], start: usize) { + if tokens.get(start).map(String::as_str) == Some("steps") { + if tokens + .get(start + 1) + .is_some_and(|token| token.parse::().is_ok()) + { + tokens[start + 1] = "*".into(); + if tokens.get(start + 2).map(String::as_str) == Some("tool_calls") + && tokens + .get(start + 3) + .is_some_and(|token| token.parse::().is_ok()) + { + tokens[start + 3] = "*".into(); + } + } + } else if tokens.get(start).map(String::as_str) == Some("subagent_trajectories") + && tokens + .get(start + 1) + .is_some_and(|token| token.parse::().is_ok()) + { + tokens[start + 1] = "*".into(); + normalize_trajectory(tokens, start + 2); + } + } + + normalize_trajectory(&mut tokens, 0); + Ok(encode_json_pointer(&tokens)) +} + +pub(crate) fn normalize_actf_pointer(source: &str, pointer: &str) -> InputResult { + let mut tokens = decode_json_pointer(pointer)?; + if source == "actf" + && tokens.first().map(String::as_str) == Some("attempts") + && tokens.get(2).map(String::as_str) == Some("trajectory") + && tokens.get(3).map(String::as_str) == Some("steps") + && tokens + .get(4) + .is_some_and(|token| token.parse::().is_ok()) + { + tokens[4] = "*".into(); + if matches!( + tokens.get(5).map(String::as_str), + Some("tools" | "observation") + ) && tokens + .get(6) + .is_some_and(|token| token.parse::().is_ok()) + { + tokens[6] = "*".into(); + } else if tokens.get(5).map(String::as_str) == Some("assistant_content") + && tokens.get(6).map(String::as_str) == Some("tool_calls") + && tokens + .get(7) + .is_some_and(|token| token.parse::().is_ok()) + { + tokens[7] = "*".into(); + } + } + Ok(encode_json_pointer(&tokens)) +} + +pub(crate) fn normalize_openai_pointer(source: &str, pointer: &str) -> InputResult { + let mut tokens = decode_json_pointer(pointer)?; + if source == "openai-msg" + && tokens.first().map(String::as_str) == Some("session_steps") + && tokens + .get(1) + .is_some_and(|token| token.parse::().is_ok()) + { + tokens[1] = "*".into(); + if tokens.get(2).map(String::as_str) == Some("messages") + && tokens + .get(3) + .is_some_and(|token| token.parse::().is_ok()) + { + tokens[3] = "*".into(); + if tokens.get(4).map(String::as_str) == Some("tool_calls") + && tokens + .get(5) + .is_some_and(|token| token.parse::().is_ok()) + { + tokens[5] = "*".into(); + } + } else if tokens.get(2).map(String::as_str) == Some("response") + && tokens.get(3).map(String::as_str) == Some("tool_calls") + && tokens + .get(4) + .is_some_and(|token| token.parse::().is_ok()) + { + tokens[4] = "*".into(); + } + } + Ok(encode_json_pointer(&tokens)) +} + +/// Normalize native AgenticMD block positions while retaining all other +/// pointer tokens literally. `blocks` is the only array in the logical +/// unknown-fields document; numeric object keys in frontmatter and header values +/// must remain distinguishable. +pub(crate) fn normalize_agenticmd_unknown_pointer( + source: &str, + pointer: &str, +) -> InputResult { + let mut tokens = decode_json_pointer(pointer)?; + if source == "agenticmd" + && tokens.first().map(String::as_str) == Some("blocks") + && tokens.get(1).is_some_and(|token| is_array_index(token)) + { + tokens[1] = "*".into(); + } + Ok(encode_json_pointer(&tokens)) +} + +fn encode_json_pointer(tokens: &[String]) -> String { + tokens.iter().fold(String::new(), |mut pointer, token| { + pointer.push('/'); + pointer.push_str(&token.replace('~', "~0").replace('/', "~1")); + pointer + }) +} + +fn is_array_index(token: &str) -> bool { + !token.is_empty() + && !(token.len() > 1 && token.starts_with('0')) + && token.parse::().is_ok() +} + +pub(crate) fn decode_json_pointer(pointer: &str) -> InputResult> { + if pointer.is_empty() { + return Ok(Vec::new()); + } + let Some(pointer) = pointer.strip_prefix('/') else { + return Err(InputIssue::invalid( + "JSON Pointer must be empty or start with '/'", + )); + }; + pointer.split('/').map(decode_pointer_token).collect() +} + +fn decode_pointer_token(token: &str) -> InputResult { + let mut decoded = String::with_capacity(token.len()); + let mut chars = token.chars(); + while let Some(character) = chars.next() { + if character != '~' { + decoded.push(character); + continue; + } + match chars.next() { + Some('0') => decoded.push('~'), + Some('1') => decoded.push('/'), + _ => { + return Err(InputIssue::invalid( + "JSON Pointer contains an invalid '~' escape", + )) + } + } + } + Ok(decoded) +} + +pub(crate) enum PointerWrite { + InsertOnly, + ReplaceSourceOwned, +} + +pub(crate) fn restore_json_pointer( + target: &mut Value, + pointer: &str, + value: Value, + write: PointerWrite, +) -> Result<()> { + let tokens = decode_json_pointer(pointer)?; + if tokens.is_empty() { + return match write { + PointerWrite::InsertOnly => anyhow::bail!( + "cannot insert unknown field at existing root JSON Pointer '{pointer}'" + ), + PointerWrite::ReplaceSourceOwned => { + *target = value; + Ok(()) + } + }; + } + + let (last, parents) = tokens + .split_last() + .expect("non-empty JSON Pointer token sequence"); + let mut parent = target; + for token in parents { + parent = match parent { + Value::Object(object) => object.get_mut(token).ok_or_else(|| { + anyhow::anyhow!("JSON Pointer '{pointer}' has a missing object parent") + })?, + Value::Array(array) => { + let index = array_index(token, pointer)?; + array.get_mut(index).ok_or_else(|| { + anyhow::anyhow!("JSON Pointer '{pointer}' references a missing array slot") + })? + } + _ => anyhow::bail!("JSON Pointer '{pointer}' has a non-container parent"), + }; + } + + match parent { + Value::Object(object) => match write { + PointerWrite::InsertOnly => { + if object.contains_key(last) { + anyhow::bail!("JSON Pointer '{pointer}' collides with a canonical value"); + } + object.insert(last.clone(), value); + } + PointerWrite::ReplaceSourceOwned => { + object.insert(last.clone(), value); + } + }, + Value::Array(array) => { + let index = array_index(last, pointer)?; + let slot = array.get_mut(index).ok_or_else(|| { + anyhow::anyhow!("JSON Pointer '{pointer}' references a missing array slot") + })?; + match write { + PointerWrite::InsertOnly => { + anyhow::bail!("JSON Pointer '{pointer}' collides with an existing array value") + } + PointerWrite::ReplaceSourceOwned => *slot = value, + } + } + _ => anyhow::bail!("JSON Pointer '{pointer}' has a non-container parent"), + } + Ok(()) +} + +fn array_index(token: &str, pointer: &str) -> Result { + if token.is_empty() || (token.len() > 1 && token.starts_with('0')) { + anyhow::bail!("JSON Pointer '{pointer}' has an invalid array index '{token}'"); + } + token.parse::().map_err(|_| { + anyhow::anyhow!("JSON Pointer '{pointer}' has an invalid array index '{token}'") + }) +} + +pub(crate) fn canonical_source_document_id(value: &Value) -> Result { + let mut value = value.clone(); + if let Value::Object(object) = &mut value { + object.remove("_storyline"); + } + let canonical = canonicalize_json_value(value); + Ok(blake3::hash(&serde_json::to_vec(&canonical)?) + .to_hex() + .to_string()) +} + +fn canonicalize_json_value(value: Value) -> Value { + match value { + Value::Array(values) => { + Value::Array(values.into_iter().map(canonicalize_json_value).collect()) + } + Value::Object(object) => { + let sorted = object + .into_iter() + .map(|(key, value)| (key, canonicalize_json_value(value))) + .collect::>(); + Value::Object(sorted.into_iter().collect()) + } + value => value, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::formats::StorylineDocument; + use serde_json::json; + + #[test] + fn envelope_distributes_foreign_sources_by_carrier() { + let mut raw = json!({ + "attempts": {"1": {}}, + "_storyline": {"unknown_fields": {"version": 1, "by_trajectory": { + "/attempts/1": {"sources": { + "atif": {"source_document_id": "a", "fields": {"/vendor": 7}} + }} + }}} + }); + let envelope = take_unknown_fields_envelope(&mut raw).unwrap(); + assert!(raw.get("_storyline").is_none()); + let mut stories = vec![StorylineDocument::new("s", "a")]; + attach_carried_unknown_fields( + envelope, + &[CarrierBinding { + story_index: 0, + pointer: "/attempts/1".into(), + }], + &mut stories, + UnknownFieldLimits::default(), + ) + .unwrap(); + assert_eq!( + stories[0].unknown_fields.sources["atif"].fields["/vendor"], + 7 + ); + } + + #[test] + fn envelope_rejects_reserved_shape_version_and_extra_keys() { + let invalid = [ + json!({"_storyline": null}), + json!({"_storyline": {"unknown_fields": {"by_trajectory": {}}}}), + json!({"_storyline": {"unknown_fields": {"version": 2, "by_trajectory": {}}}}), + json!({"_storyline": {"unknown_fields": { + "version": 1, "by_trajectory": {}, "extra": true + }}}), + json!({"_storyline": { + "unknown_fields": {"version": 1, "by_trajectory": {}}, + "extra": true + }}), + ]; + for mut raw in invalid { + assert!( + take_unknown_fields_envelope(&mut raw).is_err(), + "accepted {raw}" + ); + assert!(raw.get("_storyline").is_some()); + } + } + + #[test] + fn envelope_rejects_bad_duplicate_and_unbound_carriers() { + let mut bad_pointer = json!({"_storyline": {"unknown_fields": { + "version": 1, + "by_trajectory": {"bad": {"sources": {}}} + }}}); + assert!(take_unknown_fields_envelope(&mut bad_pointer).is_err()); + + let stories = &mut [StorylineDocument::new("s", "a")]; + assert!(attach_carried_unknown_fields( + BTreeMap::new(), + &[ + CarrierBinding { + story_index: 0, + pointer: "/same".into() + }, + CarrierBinding { + story_index: 0, + pointer: "/same".into() + }, + ], + stories, + UnknownFieldLimits::default(), + ) + .is_err()); + + let unbound = BTreeMap::from([("/missing".into(), StorylineUnknownFields::default())]); + assert!(attach_carried_unknown_fields( + unbound, + &[CarrierBinding { + story_index: 0, + pointer: "/bound".into() + }], + stories, + UnknownFieldLimits::default(), + ) + .is_err()); + } + + #[test] + fn envelope_attachment_rejects_source_id_changes_and_total_limit_splitting() { + let mut story = StorylineDocument::new("s", "a"); + story + .unknown_fields + .insert("atif", "first", "/owned", json!(1)) + .unwrap(); + story.refresh_unknown_key_counts().unwrap(); + + let changed_id = BTreeMap::from([( + "".into(), + StorylineUnknownFields { + sources: BTreeMap::from([( + "atif".into(), + SourceUnknownFields { + source_document_id: "second".into(), + fields: BTreeMap::from([("/foreign".into(), json!(2))]), + }, + )]), + }, + )]); + assert!(attach_carried_unknown_fields( + changed_id, + &[CarrierBinding { + story_index: 0, + pointer: "".into() + }], + std::slice::from_mut(&mut story), + UnknownFieldLimits::default(), + ) + .is_err()); + assert_eq!( + story.unknown_fields.sources["atif"].source_document_id, + "first" + ); + + let carried = BTreeMap::from([( + "".into(), + StorylineUnknownFields { + sources: BTreeMap::from([( + "actf".into(), + SourceUnknownFields { + source_document_id: "doc".into(), + fields: BTreeMap::from([("/foreign".into(), json!(2))]), + }, + )]), + }, + )]); + assert!(attach_carried_unknown_fields( + carried, + &[CarrierBinding { + story_index: 0, + pointer: "".into() + }], + std::slice::from_mut(&mut story), + UnknownFieldLimits { + max_fields: 1, + max_bytes: 1024 + }, + ) + .is_err()); + assert!(!story.unknown_fields.sources.contains_key("actf")); + } + + #[test] + fn envelope_writer_excludes_target_source_and_rejects_reserved_collision() { + let mut story = StorylineDocument::new("s", "a"); + story + .unknown_fields + .insert("actf", "actf-doc", "/owned", json!(1)) + .unwrap(); + story + .unknown_fields + .insert("atif", "atif-doc", "/vendor", json!(7)) + .unwrap(); + story.refresh_unknown_key_counts().unwrap(); + let carriers = [CarrierBinding { + story_index: 0, + pointer: "/attempts/1".into(), + }]; + let mut target = json!({"attempts": {"1": {}}}); + write_foreign_unknown_fields_envelope( + DocumentFormat::Actf, + &mut target, + &[story.clone()], + &carriers, + ) + .unwrap(); + assert_eq!(target["_storyline"]["unknown_fields"]["version"], 1); + let sources = + &target["_storyline"]["unknown_fields"]["by_trajectory"]["/attempts/1"]["sources"]; + assert!(sources.get("actf").is_none()); + assert_eq!(sources["atif"]["fields"]["/vendor"], 7); + + let mut collision = json!({"attempts": {"1": {}}, "_storyline": {}}); + assert!(write_foreign_unknown_fields_envelope( + DocumentFormat::Actf, + &mut collision, + &[story], + &carriers, + ) + .is_err()); + } + + // This catches a missing validation/normalization pass that would otherwise + // admit malformed pointers or return unnormalized key counts. + fn normalize_test_pointer(source: &str, pointer: &str) -> InputResult { + assert_eq!(source, "atif"); + Ok(pointer.replacen("/steps/0/", "/steps/*/", 1)) + } + + #[test] + fn unknown_fields_validate_pointer_counts_and_limits() { + let mut fields = StorylineUnknownFields::default(); + fields + .insert( + "atif", + "source-1", + "/steps/0/vendor~1field", + json!({"kept": true}), + ) + .unwrap(); + let counts = fields + .validate_with(UnknownFieldLimits::default(), normalize_test_pointer) + .unwrap(); + assert_eq!(counts["atif"]["/steps/*/vendor~1field"], 1); + + let too_many = UnknownFieldLimits { + max_fields: 0, + max_bytes: 1_048_576, + }; + assert!(fields + .validate_with(too_many, normalize_test_pointer) + .is_err()); + assert!(validate_json_pointer("/bad~2escape").is_err()); + } + + #[test] + fn unknown_field_limits_accept_exact_entry_and_byte_boundaries() { + let mut fields = StorylineUnknownFields::default(); + for index in 0..DEFAULT_MAX_UNKNOWN_FIELDS { + fields + .insert("atif", "s", format!("/{index}"), json!(null)) + .unwrap(); + } + assert!(validate_unknown_fields(&fields, UnknownFieldLimits::default()).is_ok()); + fields + .insert("atif", "s", "/too-many", json!(null)) + .unwrap(); + assert!(validate_unknown_fields(&fields, UnknownFieldLimits::default()).is_err()); + + let exact_string_len = DEFAULT_MAX_UNKNOWN_BYTES - 4; + let exact = json!("x".repeat(exact_string_len)); + let mut bytes_at_limit = StorylineUnknownFields::default(); + bytes_at_limit.insert("atif", "s", "/", exact).unwrap(); + assert!(validate_unknown_fields(&bytes_at_limit, UnknownFieldLimits::default()).is_ok()); + + let over_limit = json!("x".repeat(exact_string_len + 1)); + let mut bytes_over_limit = StorylineUnknownFields::default(); + bytes_over_limit + .insert("atif", "s", "/", over_limit) + .unwrap(); + assert!(validate_unknown_fields(&bytes_over_limit, UnknownFieldLimits::default()).is_err()); + } + + #[test] + fn insert_rejects_source_document_id_changes() { + let mut fields = StorylineUnknownFields::default(); + fields.insert("atif", "first", "/one", json!(1)).unwrap(); + assert!(fields.insert("atif", "second", "/two", json!(2)).is_err()); + } + + #[test] + fn validate_json_pointer_accepts_only_strict_rfc_6901_escapes() { + for pointer in ["", "/", "/a~0b/~1c", "/0"] { + assert!( + validate_json_pointer(pointer).is_ok(), + "rejected {pointer:?}" + ); + } + for pointer in ["not-a-pointer", "/bad~", "/bad~2", "/bad~~"] { + assert!( + validate_json_pointer(pointer).is_err(), + "accepted {pointer:?}" + ); + } + } + + #[test] + fn validation_rejects_malformed_deserialized_pointers_before_normalization() { + let fields = StorylineUnknownFields { + sources: BTreeMap::from([( + "atif".into(), + SourceUnknownFields { + source_document_id: "source".into(), + fields: BTreeMap::from([("/bad~2escape".into(), json!(true))]), + }, + )]), + }; + + assert!(fields + .validate_with(UnknownFieldLimits::default(), |_, pointer| Ok( + pointer.into() + )) + .is_err()); + } + + #[test] + fn restore_pointer_rejects_canonical_collision_and_missing_array_slot() { + let mut target = json!({"steps": [{"message": "canonical"}]}); + let error = restore_json_pointer( + &mut target, + "/steps/0/message", + json!("unknown"), + PointerWrite::InsertOnly, + ) + .unwrap_err(); + assert!(error.to_string().contains("/steps/0/message")); + + let error = restore_json_pointer( + &mut target, + "/steps/1/vendor", + json!(true), + PointerWrite::InsertOnly, + ) + .unwrap_err(); + assert!(error.to_string().contains("/steps/1/vendor")); + assert_eq!(target, json!({"steps": [{"message": "canonical"}]})); + } + + #[test] + fn canonical_source_document_id_ignores_envelope_and_object_key_order() { + let left = json!({"b": [ {"z": 1, "a": 2} ], "a": true, "_storyline": {"ignored": true}}); + let right = json!({"a": true, "b": [ {"a": 2, "z": 1} ]}); + assert_eq!( + canonical_source_document_id(&left).unwrap(), + canonical_source_document_id(&right).unwrap(), + ); + } +} diff --git a/crates/persisting-pchronicle/src/lib.rs b/crates/persisting-pchronicle/src/lib.rs index 9a9e32f6..080639cc 100644 --- a/crates/persisting-pchronicle/src/lib.rs +++ b/crates/persisting-pchronicle/src/lib.rs @@ -55,13 +55,13 @@ mod store; #[cfg(feature = "lance-store")] pub(crate) use document::{QueryCapabilities, QueryTables}; -#[cfg(any(feature = "lance-store", test))] +#[cfg(feature = "lance-store")] pub(crate) use format::DocumentFormat; +#[cfg(feature = "lance-store")] +pub(crate) use formats::storyline::StorylineAgent; pub(crate) use formats::storyline::StorylineTurn; #[cfg(any(feature = "lance-store", test))] -pub(crate) use formats::storyline::{FieldPresence, StoryLink, StorylineToolCall}; -#[cfg(feature = "lance-store")] -pub(crate) use formats::storyline::{StorylineAgent, StorylinePresence}; +pub(crate) use formats::storyline::{StoryLink, StorylineToolCall}; pub(crate) use formats::{EventIdentity, EventRecord, StorylineDocument}; pub(crate) use input::{InputIssue, InputResult}; pub type Result = anyhow::Result; diff --git a/crates/persisting-pchronicle/src/model.rs b/crates/persisting-pchronicle/src/model.rs index a73a98d1..254f1ccc 100644 --- a/crates/persisting-pchronicle/src/model.rs +++ b/crates/persisting-pchronicle/src/model.rs @@ -10,7 +10,10 @@ pub use crate::formats::llm::{ LlmToolDefinition, LlmUsage, }; pub use crate::formats::storyline::{ - FieldPresence, PresenceState, StoryLink, StorylineAgent, StorylineAgentField, - StorylineCollectionShape, StorylineDocument, StorylinePresence, StorylineRootField, - StorylineToolCall, StorylineTurn, StorylineTurnField, + StoryLink, StorylineAgent, StorylineDocument, StorylineToolCall, StorylineTurn, +}; +pub use crate::formats::unknown_fields::{ + compute_unknown_key_counts, validate_json_pointer, validate_unknown_fields, + validate_unknown_fields_with, SourceUnknownFields, StorylineUnknownFields, UnknownFieldCounts, + UnknownFieldLimits, UnknownKeyCounts, DEFAULT_MAX_UNKNOWN_BYTES, DEFAULT_MAX_UNKNOWN_FIELDS, }; diff --git a/crates/persisting-pchronicle/src/store/catalog/tests.rs b/crates/persisting-pchronicle/src/store/catalog/tests.rs index aaf0d72c..dab74a25 100644 --- a/crates/persisting-pchronicle/src/store/catalog/tests.rs +++ b/crates/persisting-pchronicle/src/store/catalog/tests.rs @@ -39,7 +39,8 @@ fn storyline(session_id: &str, run_id: &str) -> StorylineDocument { final_metrics: None, continued_trajectory_ref: None, extra: None, - presence: Default::default(), + unknown_fields: Default::default(), + unknown_key_counts: Default::default(), turns: vec![StorylineTurn { id: 1, kind: None, diff --git a/crates/persisting-pchronicle/src/store/document_source.rs b/crates/persisting-pchronicle/src/store/document_source.rs index 53db7ad5..b0329145 100644 --- a/crates/persisting-pchronicle/src/store/document_source.rs +++ b/crates/persisting-pchronicle/src/store/document_source.rs @@ -18,8 +18,9 @@ use crate::format::DocumentFormat; use crate::formats::actf::ActfDocument; use crate::formats::{parse_openai_msg_corpus_value, StorylineDocument}; +use super::files::DEFAULT_LOCAL_QUERY_MAX_RECORD_BYTES; use super::{ - datafusion_bridge::from_datafusion, AgenticMdDataSource, FileTrajectoryDataSource, + datafusion_bridge::from_datafusion, AgenticMdDataSource, AtifReader, FileTrajectoryDataSource, LocalQueryManifest, RawEventDataSource, StorylineDataSource, DEFAULT_MAX_EVENT_FALLBACK_BYTES, DEFAULT_MAX_EVENT_FALLBACK_ROWS, }; @@ -284,7 +285,16 @@ impl QueryDocumentSource for DocumentSourceImpl { late_content_materialization: false, snapshot_consistent: false, }, - DocumentFormat::OpenaiMsg | DocumentFormat::Actf => QueryCapabilities { + DocumentFormat::Actf => QueryCapabilities { + projection_pushdown: true, + filter_pushdown: FilterPushdown::Inexact, + limit_pushdown: true, + scalar_indexes: false, + streaming_decode: true, + late_content_materialization: false, + snapshot_consistent: false, + }, + DocumentFormat::OpenaiMsg => QueryCapabilities { projection_pushdown: true, filter_pushdown: FilterPushdown::Unsupported, limit_pushdown: true, @@ -326,12 +336,12 @@ where { match format { DocumentFormat::Atif => { - for file in manifest.files() { - let input = read_bounded_file(file, max_file_bytes, format)?; - let input = std::str::from_utf8(&input).context("ATIF input is not UTF-8")?; - for story in super::files::parse_atif_storylines(input)? { - on_storyline(story)?; - } + for story in AtifReader::from_manifest( + manifest, + max_file_bytes, + DEFAULT_LOCAL_QUERY_MAX_RECORD_BYTES, + ) { + on_storyline(story?)?; } } DocumentFormat::OpenaiMsg => { diff --git a/crates/persisting-pchronicle/src/store/files/actf_reader.rs b/crates/persisting-pchronicle/src/store/files/actf_reader.rs new file mode 100644 index 00000000..6a7f31f7 --- /dev/null +++ b/crates/persisting-pchronicle/src/store/files/actf_reader.rs @@ -0,0 +1,176 @@ +//! Bounded-memory ACTF document reader for local query paths. + +use std::io::{self, BufRead, Read}; +use std::path::Path; + +use anyhow::{Context, Result}; +use serde::Deserialize; + +use crate::convert::actf_to_storylines; +use crate::formats::actf::ActfDocument; +use crate::formats::storyline::StorylineDocument; +use crate::InputIssue; + +use super::json_stream::{visit_json_stream, ScopedJsonObjectReader}; + +#[cfg(test)] +pub(crate) fn parse_actf_storylines_from_reader( + path: &Path, + reader: &mut R, + max_record_bytes: usize, +) -> Result> { + parse_actf_storylines_from_reader_with_stats(path, reader, max_record_bytes) + .map(|(stories, _)| stories) +} + +pub(super) fn parse_actf_storylines_from_reader_with_stats( + path: &Path, + reader: &mut R, + max_record_bytes: usize, +) -> Result<(Vec, usize)> { + let mut stories = Vec::new(); + let visit = visit_json_stream( + path, + reader, + max_record_bytes, + &mut stories, + |reader, stories| { + push_actf_stories_from_scoped( + &mut ScopedJsonObjectReader::new(reader, max_record_bytes), + stories, + ) + }, + |record, location, stories| { + push_actf_stories_from_slice(record, stories).map_err(|error| { + io::Error::new( + error.kind(), + format!("parse ACTF {location} in {}: {error}", path.display()), + ) + }) + }, + ) + .map_err(|error| InputIssue::invalid(error.to_string()).at(path.display().to_string())) + .with_context(|| format!("read ACTF input {}", path.display()))?; + anyhow::ensure!( + visit.record_count > 0 && !stories.is_empty(), + "ACTF input contains no trajectories: {}", + path.display() + ); + Ok((stories, visit.peak_record_bytes)) +} + +fn push_actf_stories_from_scoped( + scoped: &mut ScopedJsonObjectReader<'_, R>, + stories: &mut Vec, +) -> io::Result<()> { + let document = deserialize_actf_document(&mut *scoped) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + if !scoped.is_finished() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "ACTF JSON object was not fully consumed", + )); + } + extend_actf_stories(document, stories) +} + +fn push_actf_stories_from_slice( + record: &[u8], + stories: &mut Vec, +) -> io::Result<()> { + let mut deserializer = serde_json::Deserializer::from_slice(record); + let document = ActfDocument::deserialize(&mut deserializer) + .and_then(|document| { + deserializer.end()?; + Ok(document) + }) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + document + .validate() + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error.to_string()))?; + extend_actf_stories(document, stories) +} + +fn extend_actf_stories( + document: ActfDocument, + stories: &mut Vec, +) -> io::Result<()> { + stories.extend( + actf_to_storylines(&document) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?, + ); + Ok(()) +} + +fn deserialize_actf_document(reader: R) -> Result { + let mut deserializer = serde_json::Deserializer::from_reader(reader); + let document = + ActfDocument::deserialize(&mut deserializer).context("deserialize ACTF document")?; + deserializer + .end() + .context("finish ACTF document deserialization")?; + document + .validate() + .map_err(|error| anyhow::anyhow!(error.to_string()))?; + Ok(document) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Cursor; + + fn fixture(name: &str) -> std::path::PathBuf { + std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("tests/fixtures/import_roundtrip") + .join(name) + } + + #[test] + fn streams_single_actf_object_without_whole_file_buffer() { + let path = fixture("protein-assembly_trimmed.actf.json"); + let raw = std::fs::read(&path).unwrap(); + let mut reader = Cursor::new(raw); + let stories = + parse_actf_storylines_from_reader(&path, &mut reader, 64 * 1024 * 1024).unwrap(); + assert!(!stories.is_empty()); + } + + #[test] + fn streams_actf_array_without_record_vec() { + let path = fixture("protein-assembly_trimmed.actf.json"); + let object = std::fs::read_to_string(&path).unwrap(); + let corpus = format!("[{object},{object}]"); + let mut reader = Cursor::new(corpus.into_bytes()); + let stories = + parse_actf_storylines_from_reader(&path, &mut reader, 64 * 1024 * 1024).unwrap(); + assert!(stories.len() >= 2); + } + + #[test] + fn enforces_max_record_bytes_on_actf_array_elements() { + let path = fixture("protein-assembly_trimmed.actf.json"); + let object = std::fs::read_to_string(&path).unwrap(); + let corpus = format!("[{object}]"); + let mut reader = Cursor::new(corpus.into_bytes()); + let error = parse_actf_storylines_from_reader(&path, &mut reader, 512).unwrap_err(); + assert!( + format!("{error:#}").contains("max_record_bytes 512"), + "{error:#}" + ); + } + + #[test] + fn invalid_actf_stream_preserves_the_input_issue_boundary() { + let path = Path::new("invalid.actf.json"); + let mut reader = Cursor::new(b"not-json"); + let error = parse_actf_storylines_from_reader(path, &mut reader, 1024).unwrap_err(); + + assert!( + error + .chain() + .any(|source| source.downcast_ref::().is_some()), + "missing InputIssue source: {error:#}" + ); + } +} diff --git a/crates/persisting-pchronicle/src/store/files/actf_stream.rs b/crates/persisting-pchronicle/src/store/files/actf_stream.rs new file mode 100644 index 00000000..9038c8c3 --- /dev/null +++ b/crates/persisting-pchronicle/src/store/files/actf_stream.rs @@ -0,0 +1,868 @@ +use super::json_stream::{visit_json_stream, BoundedCountingReader}; +use super::projected_steps::{ + canonical_json_text, emit_projected_step_batch, projected_timing_from_actf_metrics, + ProjectedStepRow, +}; +use super::*; +use std::fmt; +use std::io::{self, BufRead}; + +#[derive(Debug)] +struct ProjectedActfDocument { + task_id: String, + attempts: Vec<(String, ProjectedActfAttempt)>, +} + +impl ProjectedActfDocument { + fn attempt_count(&self) -> usize { + self.attempts.len() + } +} + +#[derive(Debug)] +struct ProjectedActfAttempt { + steps: Vec, +} + +impl ProjectedActfAttempt { + fn step_count(&self) -> usize { + self.steps.len() + } +} + +#[derive(Debug)] +struct ProjectedActfStep { + step_id: i64, + started_at: String, + content: String, + reasoning_content: Option, + tools_nonempty: bool, + observation_present: bool, + metrics_json: Option>, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Deserialize)] +#[serde(field_identifier, rename_all = "snake_case")] +enum ProjectedActfDocumentField { + TaskId, + Attempts, + #[serde(other)] + Other, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Deserialize)] +#[serde(field_identifier, rename_all = "snake_case")] +enum ProjectedActfAttemptField { + Trajectory, + #[serde(other)] + Other, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Deserialize)] +#[serde(field_identifier, rename_all = "snake_case")] +enum ProjectedActfTrajectoryField { + Steps, + #[serde(other)] + Other, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Deserialize)] +#[serde(field_identifier, rename_all = "snake_case")] +enum ProjectedActfStepField { + StepId, + AssistantContent, + Metric, + Tools, + Observation, + StartedAt, + #[serde(other)] + Other, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Deserialize)] +#[serde(field_identifier, rename_all = "snake_case")] +enum ProjectedActfAssistantContentField { + Content, + ReasoningContent, + #[serde(other)] + Other, +} + +struct ProjectedActfDocumentSeed<'a> { + scan: &'a FileScanSpec, +} + +impl<'de> DeserializeSeed<'de> for ProjectedActfDocumentSeed<'_> { + type Value = ProjectedActfDocument; + + fn deserialize(self, deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + deserializer.deserialize_map(ProjectedActfDocumentVisitor { scan: self.scan }) + } +} + +struct ProjectedActfDocumentVisitor<'a> { + scan: &'a FileScanSpec, +} + +impl<'de> Visitor<'de> for ProjectedActfDocumentVisitor<'_> { + type Value = ProjectedActfDocument; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("an ACTF document object") + } + + fn visit_map(self, mut map: A) -> std::result::Result + where + A: MapAccess<'de>, + { + let mut task_id = None; + let mut attempts = None; + while let Some(key) = map.next_key::()? { + match key { + ProjectedActfDocumentField::TaskId => { + task_id = Some(map.next_value::()?); + } + ProjectedActfDocumentField::Attempts => { + attempts = + Some(map.next_value_seed(ProjectedActfAttemptsSeed { scan: self.scan })?); + } + ProjectedActfDocumentField::Other => { + map.next_value::()?; + } + } + } + Ok(ProjectedActfDocument { + task_id: task_id.ok_or_else(|| de::Error::missing_field("task_id"))?, + attempts: attempts.ok_or_else(|| de::Error::missing_field("attempts"))?, + }) + } +} + +struct ProjectedActfAttemptsSeed<'a> { + scan: &'a FileScanSpec, +} + +impl<'de> DeserializeSeed<'de> for ProjectedActfAttemptsSeed<'_> { + type Value = Vec<(String, ProjectedActfAttempt)>; + + fn deserialize(self, deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + deserializer.deserialize_map(ProjectedActfAttemptsVisitor { scan: self.scan }) + } +} + +struct ProjectedActfAttemptsVisitor<'a> { + scan: &'a FileScanSpec, +} + +impl<'de> Visitor<'de> for ProjectedActfAttemptsVisitor<'_> { + type Value = Vec<(String, ProjectedActfAttempt)>; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("an ACTF attempts map") + } + + fn visit_map(self, mut map: A) -> std::result::Result + where + A: MapAccess<'de>, + { + let mut attempts = Vec::new(); + while let Some(attempt_id) = map.next_key::()? { + let attempt = map.next_value_seed(ProjectedActfAttemptSeed { scan: self.scan })?; + attempts.push((attempt_id, attempt)); + } + if attempts.is_empty() { + return Err(de::Error::custom("ACTF attempts must not be empty")); + } + Ok(attempts) + } +} + +struct ProjectedActfAttemptSeed<'a> { + scan: &'a FileScanSpec, +} + +impl<'de> DeserializeSeed<'de> for ProjectedActfAttemptSeed<'_> { + type Value = ProjectedActfAttempt; + + fn deserialize(self, deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + deserializer.deserialize_map(ProjectedActfAttemptVisitor { scan: self.scan }) + } +} + +struct ProjectedActfAttemptVisitor<'a> { + scan: &'a FileScanSpec, +} + +impl<'de> Visitor<'de> for ProjectedActfAttemptVisitor<'_> { + type Value = ProjectedActfAttempt; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("an ACTF attempt object") + } + + fn visit_map(self, mut map: A) -> std::result::Result + where + A: MapAccess<'de>, + { + let mut trajectory = None; + while let Some(key) = map.next_key::()? { + match key { + ProjectedActfAttemptField::Trajectory => { + trajectory = + Some(map.next_value_seed(ProjectedActfTrajectorySeed { scan: self.scan })?); + } + ProjectedActfAttemptField::Other => { + map.next_value::()?; + } + } + } + trajectory.ok_or_else(|| de::Error::missing_field("trajectory")) + } +} + +struct ProjectedActfTrajectorySeed<'a> { + scan: &'a FileScanSpec, +} + +impl<'de> DeserializeSeed<'de> for ProjectedActfTrajectorySeed<'_> { + type Value = ProjectedActfAttempt; + + fn deserialize(self, deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + deserializer.deserialize_map(ProjectedActfTrajectoryVisitor { scan: self.scan }) + } +} + +struct ProjectedActfTrajectoryVisitor<'a> { + scan: &'a FileScanSpec, +} + +impl<'de> Visitor<'de> for ProjectedActfTrajectoryVisitor<'_> { + type Value = ProjectedActfAttempt; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("an ACTF trajectory object") + } + + fn visit_map(self, mut map: A) -> std::result::Result + where + A: MapAccess<'de>, + { + let mut steps = None; + while let Some(key) = map.next_key::()? { + match key { + ProjectedActfTrajectoryField::Steps => { + steps = Some(map.next_value_seed(ProjectedActfStepsSeed { scan: self.scan })?); + } + ProjectedActfTrajectoryField::Other => { + map.next_value::()?; + } + } + } + steps.ok_or_else(|| de::Error::missing_field("steps")) + } +} + +struct ProjectedActfStepsSeed<'a> { + scan: &'a FileScanSpec, +} + +impl<'de> DeserializeSeed<'de> for ProjectedActfStepsSeed<'_> { + type Value = ProjectedActfAttempt; + + fn deserialize(self, deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + deserializer.deserialize_seq(ProjectedActfStepsVisitor { scan: self.scan }) + } +} + +struct ProjectedActfStepsVisitor<'a> { + scan: &'a FileScanSpec, +} + +impl<'de> Visitor<'de> for ProjectedActfStepsVisitor<'_> { + type Value = ProjectedActfAttempt; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("an ACTF steps array") + } + + fn visit_seq(self, mut sequence: A) -> std::result::Result + where + A: SeqAccess<'de>, + { + let mut steps = Vec::new(); + while let Some(step) = + sequence.next_element_seed(ProjectedActfStepSeed { scan: self.scan })? + { + steps.push(step); + } + if steps.is_empty() { + return Err(de::Error::custom("ACTF trajectory steps must not be empty")); + } + Ok(ProjectedActfAttempt { steps }) + } +} + +struct ProjectedActfStepSeed<'a> { + scan: &'a FileScanSpec, +} + +impl<'de> DeserializeSeed<'de> for ProjectedActfStepSeed<'_> { + type Value = ProjectedActfStep; + + fn deserialize(self, deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + deserializer.deserialize_map(ProjectedActfStepVisitor { scan: self.scan }) + } +} + +struct ProjectedActfStepVisitor<'a> { + scan: &'a FileScanSpec, +} + +impl<'de> Visitor<'de> for ProjectedActfStepVisitor<'_> { + type Value = ProjectedActfStep; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("an ACTF step object") + } + + fn visit_map(self, mut map: A) -> std::result::Result + where + A: MapAccess<'de>, + { + let mut step_id = None; + let mut started_at = None; + let mut content = String::new(); + let mut reasoning_content = None; + let mut tools_nonempty = false; + let mut observation_present = false; + let mut metrics_json = None; + + while let Some(key) = map.next_key::()? { + match key { + ProjectedActfStepField::StepId => { + step_id = Some(map.next_value::()?); + } + ProjectedActfStepField::StartedAt => { + started_at = Some(map.next_value::()?); + } + ProjectedActfStepField::AssistantContent => { + let assistant = + map.next_value_seed(ProjectedActfAssistantContentSeed { scan: self.scan })?; + content = assistant.content; + reasoning_content = assistant.reasoning_content; + } + ProjectedActfStepField::Metric => { + if self.scan.wants("metrics_json") || self.scan.wants("latency_ms") { + metrics_json = + map.next_value::>>()?; + } else { + map.next_value::()?; + } + } + ProjectedActfStepField::Tools => { + if self.scan.wants("kind") || self.scan.wants("effective_kind") { + tools_nonempty = map + .next_value::>>()? + .is_some_and(|calls| !calls.is_empty()); + } else { + map.next_value::()?; + } + } + ProjectedActfStepField::Observation => { + if self.scan.wants("had_observation") { + observation_present = map + .next_value::>>()? + .is_some_and(|observations| !observations.is_empty()); + } else { + map.next_value::()?; + } + } + ProjectedActfStepField::Other => { + map.next_value::()?; + } + } + } + + Ok(ProjectedActfStep { + step_id: step_id.ok_or_else(|| de::Error::missing_field("step_id"))?, + started_at: started_at.ok_or_else(|| de::Error::missing_field("started_at"))?, + content, + reasoning_content, + tools_nonempty, + observation_present, + metrics_json, + }) + } +} + +struct ProjectedActfAssistantContent { + content: String, + reasoning_content: Option, +} + +struct ProjectedActfAssistantContentSeed<'a> { + scan: &'a FileScanSpec, +} + +impl<'de> DeserializeSeed<'de> for ProjectedActfAssistantContentSeed<'_> { + type Value = ProjectedActfAssistantContent; + + fn deserialize(self, deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + deserializer.deserialize_map(ProjectedActfAssistantContentVisitor { scan: self.scan }) + } +} + +struct ProjectedActfAssistantContentVisitor<'a> { + scan: &'a FileScanSpec, +} + +impl<'de> Visitor<'de> for ProjectedActfAssistantContentVisitor<'_> { + type Value = ProjectedActfAssistantContent; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("ACTF assistant_content object") + } + + fn visit_map(self, mut map: A) -> std::result::Result + where + A: MapAccess<'de>, + { + let mut content = String::new(); + let mut reasoning_content = None; + while let Some(key) = map.next_key::()? { + match key { + ProjectedActfAssistantContentField::Content => { + if self.scan.wants("message_json") { + content = map.next_value::()?; + } else { + map.next_value::()?; + } + } + ProjectedActfAssistantContentField::ReasoningContent => { + if self.scan.wants("reasoning_content") { + let value = map.next_value::()?; + if !value.is_empty() { + reasoning_content = Some(value); + } + } else { + map.next_value::()?; + } + } + ProjectedActfAssistantContentField::Other => { + map.next_value::()?; + } + } + } + Ok(ProjectedActfAssistantContent { + content, + reasoning_content, + }) + } +} + +const PROJECTED_QUERY_CANCELLED: &str = "pChronicle projected query receiver closed"; + +struct ProjectedActfStream<'a> { + file: &'a Arc, + runtime: &'a Arc, + schema: &'a SchemaRef, + batch_size: usize, + scan: &'a FileScanSpec, + tx: &'a Sender>, + pending: Vec, + document_ids: HashSet, + cancelled: bool, +} + +impl<'a> ProjectedActfStream<'a> { + fn new( + file: &'a Arc, + runtime: &'a Arc, + schema: &'a SchemaRef, + batch_size: usize, + scan: &'a FileScanSpec, + tx: &'a Sender>, + ) -> Self { + Self { + file, + runtime, + schema, + batch_size, + scan, + tx, + pending: Vec::with_capacity(batch_size), + document_ids: HashSet::new(), + cancelled: false, + } + } + + fn consume(&mut self, document: ProjectedActfDocument) -> Result<()> { + self.runtime + .metrics + .inner + .streamed_records + .fetch_add(1, Ordering::Relaxed); + if !project_actf_document( + document, + self.file, + self.runtime, + self.schema, + self.batch_size, + self.scan, + self.tx, + &mut self.pending, + &mut self.document_ids, + )? { + self.cancelled = true; + anyhow::bail!(PROJECTED_QUERY_CANCELLED); + } + Ok(()) + } + + fn finish(&mut self) -> Result<()> { + if !self.cancelled { + let _ = emit_projected_step_batch( + &mut self.pending, + self.file, + self.runtime, + self.schema, + self.tx, + )?; + } + Ok(()) + } +} + +fn deserialize_projected_actf_from_slice( + record: &[u8], + scan: &FileScanSpec, +) -> Result { + let mut deserializer = serde_json::Deserializer::from_slice(record); + let document = ProjectedActfDocumentSeed { scan } + .deserialize(&mut deserializer) + .map_err(anyhow::Error::from)?; + deserializer.end().map_err(anyhow::Error::from)?; + Ok(document) +} + +fn consume_projected_actf_reader( + reader: &mut R, + scan: &FileScanSpec, + stream: &mut ProjectedActfStream<'_>, +) -> Result<()> { + let mut deserializer = serde_json::Deserializer::from_reader(reader); + let document = ProjectedActfDocumentSeed { scan } + .deserialize(&mut deserializer) + .map_err(anyhow::Error::from)?; + deserializer.end().map_err(anyhow::Error::from)?; + stream.consume(document) +} + +pub(super) fn stream_projected_actf_steps( + file: &Arc, + runtime: &Arc, + schema: &SchemaRef, + batch_size: usize, + scan: &FileScanSpec, + tx: &Sender>, +) -> Result<()> { + let _permit = runtime.limiter.acquire()?; + file.file.validate_unchanged()?; + anyhow::ensure!( + file.file.size_bytes() <= runtime.options.max_file_bytes, + "ACTF input {} is {} bytes, exceeding max_file_bytes {}", + file.file.path().display(), + file.file.size_bytes(), + runtime.options.max_file_bytes + ); + let input = File::open(file.file.path()) + .with_context(|| format!("open ACTF input {}", file.file.path().display()))?; + let mut reader = BufReader::with_capacity( + 64 * 1024, + BoundedCountingReader::new(input, runtime.options.max_file_bytes), + ); + runtime + .metrics + .inner + .streaming_buffer_peak_bytes + .fetch_max(reader.capacity() as u64, Ordering::Relaxed); + let mut stream = ProjectedActfStream::new(file, runtime, schema, batch_size, scan, tx); + + let reader_capacity = reader.capacity() as u64; + let max_record_bytes = runtime.options.max_record_bytes; + let result = visit_json_stream( + file.file.path(), + &mut reader, + max_record_bytes, + &mut stream, + // A single object spans the whole file and is already bounded by + // `max_file_bytes`, so deserialize it without an intermediate copy. + |reader, stream| { + consume_projected_actf_reader(reader, scan, stream) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error)) + }, + |record, location, stream| { + runtime.metrics.inner.streaming_buffer_peak_bytes.fetch_max( + reader_capacity.saturating_add(record.len() as u64), + Ordering::Relaxed, + ); + let document = deserialize_projected_actf_from_slice(record, scan) + .with_context(|| { + format!( + "parse projected ACTF {location} in {}", + file.file.path().display() + ) + }) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + stream + .consume(document) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error)) + }, + ) + .map(|_| ()) + .map_err(anyhow::Error::from) + .or_else(|error| { + if stream.cancelled { + Ok(()) + } else if error.to_string().contains("JSON array contains no objects") { + Err(anyhow::anyhow!("ACTF input contains no trajectories")) + } else { + Err(error) + } + }); + if let Err(error) = result { + if !stream.cancelled { + return Err(error).with_context(|| { + format!("parse projected ACTF input {}", file.file.path().display()) + }); + } + } + stream.finish()?; + let bytes_read = reader.get_ref().bytes_read(); + runtime + .metrics + .inner + .source_bytes_read + .fetch_add(bytes_read, Ordering::Relaxed); + runtime + .metrics + .inner + .files_parsed + .fetch_add(1, Ordering::Relaxed); + runtime + .metrics + .inner + .projected_files + .fetch_add(1, Ordering::Relaxed); + if !stream.cancelled { + file.file.validate_unchanged()?; + } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +fn project_actf_document( + document: ProjectedActfDocument, + file: &Arc, + runtime: &Arc, + schema: &SchemaRef, + batch_size: usize, + scan: &FileScanSpec, + tx: &Sender>, + pending: &mut Vec, + document_ids: &mut HashSet, +) -> Result { + anyhow::ensure!( + !document.task_id.trim().is_empty(), + "ACTF task_id is required" + ); + let multiple_attempts = document.attempt_count() > 1; + for (attempt_id, attempt) in document.attempts { + let session_id = if multiple_attempts { + format!("{}#attempt-{attempt_id}", document.task_id) + } else { + document.task_id.clone() + }; + let document_id = session_id.clone(); + anyhow::ensure!( + document_ids.insert(document_id.clone()), + "duplicate ACTF document_id '{}' in {}", + document_id, + file.file.path().display() + ); + runtime + .metrics + .inner + .documents_scanned + .fetch_add(1, Ordering::Relaxed); + runtime + .metrics + .inner + .rows_scanned + .fetch_add(attempt.step_count() as u64, Ordering::Relaxed); + if !scan.matches_document(&session_id) { + runtime + .metrics + .inner + .documents_pruned + .fetch_add(1, Ordering::Relaxed); + runtime + .metrics + .inner + .rows_pruned + .fetch_add(attempt.step_count() as u64, Ordering::Relaxed); + continue; + } + + let mut rows = Vec::with_capacity(attempt.steps.len()); + let mut step_ids = HashSet::with_capacity(attempt.steps.len()); + for step in attempt.steps { + anyhow::ensure!(step.step_id >= 1, "ACTF step_id must start from 1"); + anyhow::ensure!( + step_ids.insert(step.step_id), + "duplicate ACTF step_id {} in document {}", + step.step_id, + document_id + ); + if !scan.matches_step(step.step_id, "agent") { + runtime + .metrics + .inner + .rows_pruned + .fetch_add(1, Ordering::Relaxed); + continue; + } + rows.push(project_actf_step( + &document_id, + &session_id, + &document.task_id, + step, + scan, + )); + } + rows.sort_by_key(|row| row.step_id); + runtime + .metrics + .inner + .rows_emitted + .fetch_add(rows.len() as u64, Ordering::Relaxed); + for row in rows { + pending.push(row); + if pending.len() == batch_size + && !emit_projected_step_batch(pending, file, runtime, schema, tx)? + { + return Ok(false); + } + } + } + Ok(true) +} + +fn project_actf_step( + document_id: &str, + session_id: &str, + task_id: &str, + step: ProjectedActfStep, + scan: &FileScanSpec, +) -> ProjectedStepRow { + let effective_kind = if step.tools_nonempty { + "autonomous" + } else { + "dialogue" + }; + let kind = step.tools_nonempty.then(|| "autonomous".to_string()); + let latency_ms = + projected_timing_from_actf_metrics(step.metrics_json.as_ref().map(|value| value.get())); + let message_json = if scan.wants("message_json") { + serde_json::to_string(&step.content).unwrap_or_else(|_| "\"\"".to_string()) + } else { + "null".to_string() + }; + + ProjectedStepRow { + document_id: if scan.wants("document_id") { + document_id.to_string() + } else { + String::new() + }, + run_id: scan.wants("run_id").then(|| task_id.to_string()), + session_id: if scan.wants("session_id") { + session_id.to_string() + } else { + String::new() + }, + step_id: step.step_id, + kind: scan.wants("kind").then_some(kind).flatten(), + effective_kind: if scan.wants("effective_kind") { + effective_kind.to_string() + } else { + String::new() + }, + timestamp: scan.wants("timestamp").then_some(step.started_at), + source: if scan.wants("source") { + "agent".to_string() + } else { + String::new() + }, + message_json, + reasoning_content: scan + .wants("reasoning_content") + .then_some(step.reasoning_content) + .flatten(), + reasoning_effort_json: None, + metrics_json: scan + .wants("metrics_json") + .then_some(step.metrics_json.as_deref().map(canonical_json_text)) + .flatten(), + model_name: None, + llm_call_count: scan.wants("llm_call_count").then_some(1), + is_copied_context: None, + latency_ms: scan.wants("latency_ms").then_some(latency_ms).flatten(), + ttft_ms: None, + had_observation: scan.wants("had_observation") && step.observation_present, + extra_json: None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn fixture(name: &str) -> std::path::PathBuf { + std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("tests/fixtures/import_roundtrip") + .join(name) + } + + #[test] + fn projected_actf_document_parses_trimmed_fixture() { + let path = fixture("protein-assembly_trimmed.actf.json"); + let raw = std::fs::read(&path).unwrap(); + let scan = FileScanSpec::new(Some(&vec![0, 1, 2, 3]), &[], &story_steps_arrow_schema()); + let document = deserialize_projected_actf_from_slice(&raw, &scan).unwrap(); + assert_eq!(document.task_id, "protein-assembly-trimmed"); + assert_eq!(document.attempts.len(), 1); + assert_eq!(document.attempts[0].1.steps.len(), 2); + } +} diff --git a/crates/persisting-pchronicle/src/store/files/atif_reader.rs b/crates/persisting-pchronicle/src/store/files/atif_reader.rs index 38cb9bc4..9f3fa21f 100644 --- a/crates/persisting-pchronicle/src/store/files/atif_reader.rs +++ b/crates/persisting-pchronicle/src/store/files/atif_reader.rs @@ -1,75 +1,100 @@ //! Bounded-memory ATIF document reader shared by conversion and query paths. use std::fs::File; -use std::io::{BufRead, BufReader, Lines}; -use std::path::{Path, PathBuf}; +use std::io::{self, BufRead, BufReader, Read}; +use std::path::Path; use anyhow::{Context, Result}; +use serde::Deserialize; -use crate::atif::AtifTrajectory; use crate::convert::atif_collection_to_storylines; use crate::format::DocumentFormat; -use crate::formats::storyline::{StorylineCollectionShape, StorylineDocument}; +use crate::formats::storyline::StorylineDocument; +use crate::InputIssue; -use super::{LocalQueryInputFile, LocalQueryManifest}; +use super::json_stream::{ + read_bounded_line, trim_ascii_whitespace, visit_json_stream, BoundedCountingReader, + ScopedJsonObjectReader, +}; +use super::{ + LocalQueryInputFile, LocalQueryManifest, DEFAULT_LOCAL_QUERY_MAX_FILE_BYTES, + DEFAULT_LOCAL_QUERY_MAX_RECORD_BYTES, +}; /// Bounded-memory ATIF reader. /// -/// JSONL/NDJSON inputs are decoded one non-empty line at a time. Directories -/// are traversed in stable path order and only the current file is open. A -/// regular `.json` file may contain one object or an array and is buffered per -/// file for compatibility; large corpora should use NDJSON. -pub struct AtifReader { - files: std::vec::IntoIter, +/// JSONL/NDJSON inputs are decoded one non-empty line at a time. Object and +/// array `.json` files are streamed through the shared JSON document reader +/// without loading the whole file into memory first. +pub(crate) struct AtifReader { + files: std::vec::IntoIter, current: Option, pending: std::vec::IntoIter, + max_file_bytes: u64, + max_record_bytes: usize, } enum AtifFileReader { - Lines { - path: PathBuf, - lines: Lines>, + Ndjson { + file: LocalQueryInputFile, + reader: BufReader>, + record: Vec, line_number: usize, - root_ordinal: i64, }, Documents(std::vec::IntoIter), } impl AtifReader { - pub fn open(path: impl AsRef) -> Result { + pub(crate) fn open(path: impl AsRef) -> Result { let manifest = LocalQueryManifest::for_format(path, DocumentFormat::Atif)?; - Ok(Self::from_files(manifest.files())) + Ok(Self::from_manifest( + &manifest, + DEFAULT_LOCAL_QUERY_MAX_FILE_BYTES, + DEFAULT_LOCAL_QUERY_MAX_RECORD_BYTES, + )) } - fn from_files(files: &[LocalQueryInputFile]) -> Self { + pub(crate) fn from_manifest( + manifest: &LocalQueryManifest, + max_file_bytes: u64, + max_record_bytes: usize, + ) -> Self { Self { - files: files - .iter() - .map(|file| file.path().to_path_buf()) - .collect::>() - .into_iter(), + files: manifest.files().to_vec().into_iter(), current: None, pending: Vec::new().into_iter(), + max_file_bytes, + max_record_bytes, } } - fn open_file(path: PathBuf) -> Result { - match path.extension().and_then(|value| value.to_str()) { - Some("jsonl" | "ndjson") => { - let file = File::open(&path) - .with_context(|| format!("open ATIF datasource {}", path.display()))?; - Ok(AtifFileReader::Lines { - path, - lines: BufReader::new(file).lines(), - line_number: 0, - root_ordinal: 0, - }) - } + fn open_file(&self, file: LocalQueryInputFile) -> Result { + file.validate_unchanged()?; + anyhow::ensure!( + file.size_bytes() <= self.max_file_bytes, + "ATIF input {} is {} bytes, exceeding max_file_bytes {}", + file.path().display(), + file.size_bytes(), + self.max_file_bytes + ); + let input = File::open(file.path()) + .with_context(|| format!("open ATIF datasource {}", file.path().display()))?; + let mut reader = BufReader::new(BoundedCountingReader::new(input, self.max_file_bytes)); + match file.path().extension().and_then(|value| value.to_str()) { + Some("jsonl" | "ndjson") => Ok(AtifFileReader::Ndjson { + file, + reader, + record: Vec::new(), + line_number: 0, + }), _ => { - let input = std::fs::read_to_string(&path) - .with_context(|| format!("read ATIF datasource {}", path.display()))?; - let documents = parse_storylines(&input) - .with_context(|| format!("parse ATIF datasource {}", path.display()))?; + let documents = parse_atif_storylines_from_reader( + file.path(), + &mut reader, + self.max_record_bytes, + ) + .with_context(|| format!("parse ATIF datasource {}", file.path().display()))?; + file.validate_unchanged()?; Ok(AtifFileReader::Documents(documents.into_iter())) } } @@ -91,68 +116,75 @@ impl Iterator for AtifReader { return Some(Ok(document)); } } - AtifFileReader::Lines { - path, - lines, + AtifFileReader::Ndjson { + file, + reader, + record, line_number, - root_ordinal, - } => { - for line in lines.by_ref() { - *line_number += 1; - let line = match line { - Ok(line) => line, - Err(error) => { - return Some(Err(error).with_context(|| { + } => loop { + *line_number += 1; + let length = match read_bounded_line(reader, record, self.max_record_bytes) + { + Ok(length) => length, + Err(error) => { + return Some( + Err(anyhow::Error::new( + InputIssue::invalid(error.to_string()).at(format!( + "{} line {}", + file.path().display(), + line_number + )), + )) + .with_context(|| { format!( "read ATIF datasource {} line {}", - path.display(), + file.path().display(), line_number ) - })); - } - }; - if line.trim().is_empty() { - continue; + }), + ); } - let trajectory = match AtifTrajectory::from_json_str(line.trim()) - .map_err(anyhow::Error::from) + }; + if length == 0 { + if let Err(error) = file.validate_unchanged() { + return Some(Err(error)); + } + break; + } + let line = trim_ascii_whitespace(record); + if line.is_empty() { + continue; + } + let value = + match serde_json::from_slice(line) + .map_err(|error| { + anyhow::Error::new(InputIssue::invalid(error.to_string()).at( + format!("{} line {}", file.path().display(), line_number), + )) + }) .with_context(|| { format!( "parse ATIF datasource {} line {}", - path.display(), + file.path().display(), line_number ) }) { - Ok(trajectory) => trajectory, + Ok(value) => value, Err(error) => return Some(Err(error)), }; - let ordinal = *root_ordinal; - *root_ordinal = match root_ordinal.checked_add(1) { - Some(next) => next, - None => { - return Some(Err(anyhow::anyhow!( - "ATIF collection ordinal overflow" - ))) - } - }; - let stories = match atif_collection_to_storylines( - &trajectory, - StorylineCollectionShape::Sequence, - ordinal, - ) { - Ok(stories) => stories, - Err(error) => return Some(Err(error)), - }; - self.pending = stories.into_iter(); - return self.pending.next().map(Ok); - } - } + let stories = match atif_collection_to_storylines(value) { + Ok(stories) => stories, + Err(error) => return Some(Err(error)), + }; + self.pending = stories.into_iter(); + return self.pending.next().map(Ok); + }, } self.current = None; } - let path = self.files.next()?; - match Self::open_file(path) { + let file = self.files.next()?; + match self.open_file(file) { Ok(reader) => self.current = Some(reader), Err(error) => return Some(Err(error)), } @@ -160,74 +192,147 @@ impl Iterator for AtifReader { } } -#[cfg(test)] -fn parse_documents(input: &str) -> Result> { - parse_documents_with_shape(input).map(|(_, documents)| documents) +fn parse_atif_storylines_from_reader( + path: &Path, + reader: &mut R, + max_record_bytes: usize, +) -> Result> { + parse_atif_storylines_from_reader_with_stats(path, reader, max_record_bytes) + .map(|(stories, _)| stories) } -pub(crate) fn parse_storylines(input: &str) -> Result> { - let (shape, documents) = parse_documents_with_shape(input)?; +pub(super) fn parse_atif_storylines_from_reader_with_stats( + path: &Path, + reader: &mut R, + max_record_bytes: usize, +) -> Result<(Vec, usize)> { let mut stories = Vec::new(); - for (ordinal, trajectory) in documents.into_iter().enumerate() { - let ordinal = i64::try_from(ordinal).context("ATIF collection ordinal overflow")?; - stories.extend(atif_collection_to_storylines(&trajectory, shape, ordinal)?); - } - Ok(stories) + let visit = visit_json_stream( + path, + reader, + max_record_bytes, + &mut stories, + |reader, stories| { + push_atif_stories_from_scoped( + &mut ScopedJsonObjectReader::new(reader, max_record_bytes), + stories, + ) + }, + |record, location, stories| { + push_atif_stories_from_slice(record, stories).map_err(|error| { + io::Error::new( + error.kind(), + format!("parse ATIF {location} in {}: {error}", path.display()), + ) + }) + }, + ) + .map_err(|error| InputIssue::invalid(error.to_string()).at(path.display().to_string())) + .with_context(|| format!("read ATIF input {}", path.display()))?; + anyhow::ensure!( + visit.record_count > 0 && !stories.is_empty(), + "ATIF input contains no trajectories: {}", + path.display() + ); + Ok((stories, visit.peak_record_bytes)) } -fn parse_documents_with_shape( - input: &str, -) -> Result<(StorylineCollectionShape, Vec)> { - let trimmed = input.trim(); - if trimmed.is_empty() { - anyhow::bail!("ATIF input is empty"); - } - if let Ok(trajectory) = serde_json::from_str::(trimmed) { - trajectory.validate().map_err(anyhow::Error::from)?; - return Ok((StorylineCollectionShape::Single, vec![trajectory])); - } - if let Ok(trajectories) = serde_json::from_str::>(trimmed) { - anyhow::ensure!( - !trajectories.is_empty(), - "ATIF input contains no trajectories" - ); - for trajectory in &trajectories { - trajectory.validate().map_err(anyhow::Error::from)?; - } - return Ok((StorylineCollectionShape::Sequence, trajectories)); - } - let trajectories = trimmed - .lines() - .enumerate() - .filter(|(_, line)| !line.trim().is_empty()) - .map(|(index, line)| { - AtifTrajectory::from_json_str(line) - .map_err(anyhow::Error::from) - .with_context(|| format!("parse ATIF JSONL line {}", index + 1)) +fn push_atif_stories_from_slice( + record: &[u8], + stories: &mut Vec, +) -> io::Result<()> { + let mut deserializer = serde_json::Deserializer::from_slice(record); + let value = serde_json::Value::deserialize(&mut deserializer) + .and_then(|value| { + deserializer.end()?; + Ok(value) }) - .collect::>>()?; - anyhow::ensure!( - !trajectories.is_empty(), - "ATIF input contains no trajectories" + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + extend_atif_stories(value, stories) +} + +fn push_atif_stories_from_scoped( + scoped: &mut ScopedJsonObjectReader<'_, R>, + stories: &mut Vec, +) -> io::Result<()> { + let value = deserialize_atif_value(&mut *scoped) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + if !scoped.is_finished() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "ATIF JSON object was not fully consumed", + )); + } + extend_atif_stories(value, stories) +} + +fn extend_atif_stories( + value: serde_json::Value, + stories: &mut Vec, +) -> io::Result<()> { + stories.extend( + atif_collection_to_storylines(value) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?, ); - Ok((StorylineCollectionShape::Sequence, trajectories)) + Ok(()) +} + +fn deserialize_atif_value(reader: R) -> Result { + let mut deserializer = serde_json::Deserializer::from_reader(reader); + let value = + serde_json::Value::deserialize(&mut deserializer).context("deserialize ATIF trajectory")?; + deserializer + .end() + .context("finish ATIF trajectory deserialization")?; + Ok(value) } #[cfg(test)] mod tests { use super::*; + use std::io::Cursor; + use std::path::PathBuf; + + fn atif_fixture(name: &str) -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("tests/fixtures/atif") + .join(name) + } #[test] - fn parses_object_array_and_jsonl() { - let object = r#"{"schema_version":"ATIF-v1.7","session_id":"s","agent":{"name":"a","version":"1"},"steps":[]}"#; - assert_eq!(parse_documents(object).unwrap().len(), 1); - assert_eq!(parse_documents(&format!("[{object}]")).unwrap().len(), 1); - assert_eq!( - parse_documents(&format!("{object}\n{object}\n")) - .unwrap() - .len(), - 2 + fn streams_atif_fixture_without_whole_file_buffer() { + let path = atif_fixture("dialogue_10.json"); + let raw = std::fs::read(&path).unwrap(); + let mut reader = Cursor::new(raw); + let stories = parse_atif_storylines_from_reader( + &path, + &mut reader, + DEFAULT_LOCAL_QUERY_MAX_RECORD_BYTES, + ) + .unwrap(); + assert!(!stories.is_empty()); + } + + #[test] + fn ndjson_reader_enforces_the_configured_record_limit() { + let input = tempfile::NamedTempFile::with_suffix(".ndjson").unwrap(); + let trajectory: serde_json::Value = serde_json::from_str( + &std::fs::read_to_string(atif_fixture("dialogue_10.json")).unwrap(), + ) + .unwrap(); + std::fs::write( + input.path(), + format!("{}\n", serde_json::to_string(&trajectory).unwrap()), + ) + .unwrap(); + let manifest = LocalQueryManifest::for_format(input.path(), DocumentFormat::Atif).unwrap(); + let mut reader = + AtifReader::from_manifest(&manifest, DEFAULT_LOCAL_QUERY_MAX_FILE_BYTES, 512); + + let error = reader.next().unwrap().unwrap_err(); + assert!( + format!("{error:#}").contains("max_record_bytes 512"), + "{error:#}" ); - assert!(parse_documents("").is_err()); } } diff --git a/crates/persisting-pchronicle/src/store/files/atif_stream.rs b/crates/persisting-pchronicle/src/store/files/atif_stream.rs index c556741d..7d4f0ee3 100644 --- a/crates/persisting-pchronicle/src/store/files/atif_stream.rs +++ b/crates/persisting-pchronicle/src/store/files/atif_stream.rs @@ -1,4 +1,119 @@ +use super::json_stream::{visit_json_stream, BoundedCountingReader}; +use super::projected_steps::{ + canonical_json_text, emit_projected_step_batch, projected_timing_from_metrics, ProjectedStepRow, +}; use super::*; +use serde::Deserialize; +use std::io::{self, BufRead}; + +#[derive(Clone, Copy)] +struct ProjectedAtifScanFlags { + timestamp: bool, + model_name: bool, + reasoning_effort_json: bool, + message_json: bool, + reasoning_content: bool, + kind_fields: bool, + had_observation: bool, + metrics: bool, + extra_json: bool, + llm_call_count: bool, + is_copied_context: bool, +} + +impl ProjectedAtifScanFlags { + fn new(scan: &FileScanSpec) -> Self { + Self { + timestamp: scan.wants("timestamp"), + model_name: scan.wants("model_name"), + reasoning_effort_json: scan.wants("reasoning_effort_json"), + message_json: scan.wants("message_json"), + reasoning_content: scan.wants("reasoning_content"), + kind_fields: scan.wants("kind") || scan.wants("effective_kind"), + had_observation: scan.wants("had_observation"), + metrics: scan.wants("metrics_json") + || scan.wants("latency_ms") + || scan.wants("ttft_ms"), + extra_json: scan.wants("extra_json"), + llm_call_count: scan.wants("llm_call_count"), + is_copied_context: scan.wants("is_copied_context"), + } + } + + fn source_only_steps(self) -> bool { + !self.timestamp + && !self.model_name + && !self.reasoning_effort_json + && !self.message_json + && !self.reasoning_content + && !self.kind_fields + && !self.had_observation + && !self.metrics + && !self.extra_json + && !self.llm_call_count + && !self.is_copied_context + } +} + +fn raw_json_value_present(raw: &str) -> bool { + !matches!(raw.trim(), "" | "null" | "[]" | "{}") +} + +#[derive(Debug, serde::Deserialize)] +struct SourceOnlyAtifStep { + step_id: i64, + source: String, +} + +#[derive(Debug, serde::Deserialize)] +struct SourceOnlyAtifTrajectory { + schema_version: String, + session_id: Option, + trajectory_id: Option, + agent: ProjectedAtifAgent, + steps: Vec, + #[serde(default)] + subagent_trajectories: Option>, +} + +impl From for ProjectedAtifStep { + fn from(step: SourceOnlyAtifStep) -> Self { + Self { + step_id: step.step_id, + source: step.source, + timestamp: None, + model_name: None, + reasoning_effort_json: None, + message_json: None, + reasoning_content: None, + tool_calls_nonempty: false, + observation_present: false, + metrics_json: None, + extra_json: None, + llm_call_count: None, + is_copied_context: None, + } + } +} + +impl From for ProjectedAtifTrajectory { + fn from(trajectory: SourceOnlyAtifTrajectory) -> Self { + Self { + schema_version: trajectory.schema_version, + session_id: trajectory.session_id, + trajectory_id: trajectory.trajectory_id, + agent: trajectory.agent, + steps: trajectory.steps.into_iter().map(Into::into).collect(), + subagent_trajectories: trajectory + .subagent_trajectories + .unwrap_or_default() + .into_iter() + .map(Into::into) + .collect(), + skipped_steps: 0, + } + } +} #[derive(Debug, serde::Deserialize)] struct ProjectedAtifAgent { @@ -45,13 +160,13 @@ struct ProjectedAtifStep { timestamp: Option, source: String, model_name: Option, - reasoning_effort: Option, - message: serde_json::Value, + reasoning_effort_json: Option>, + message_json: Option>, reasoning_content: Option, tool_calls_nonempty: bool, observation_present: bool, - metrics: Option, - extra: Option, + metrics_json: Option>, + extra_json: Option>, llm_call_count: Option, is_copied_context: Option, } @@ -103,27 +218,6 @@ enum ProjectedAtifStepField { Other, } -impl ProjectedAtifStepField { - fn name(self) -> &'static str { - match self { - Self::StepId => "step_id", - Self::Timestamp => "timestamp", - Self::Source => "source", - Self::ModelName => "model_name", - Self::ReasoningEffort => "reasoning_effort", - Self::Message => "message", - Self::ReasoningContent => "reasoning_content", - Self::ToolCalls => "tool_calls", - Self::Observation => "observation", - Self::Metrics => "metrics", - Self::Extra => "extra", - Self::LlmCallCount => "llm_call_count", - Self::IsCopiedContext => "is_copied_context", - Self::Other => "", - } - } -} - struct ProjectedAtifTrajectorySeed<'a> { scan: &'a FileScanSpec, } @@ -135,7 +229,13 @@ impl<'de> DeserializeSeed<'de> for ProjectedAtifTrajectorySeed<'_> { where D: serde::Deserializer<'de>, { - deserializer.deserialize_map(ProjectedAtifTrajectoryVisitor { scan: self.scan }) + if ProjectedAtifScanFlags::new(self.scan).source_only_steps() + && self.scan.step_filters.is_empty() + { + SourceOnlyAtifTrajectory::deserialize(deserializer).map(Into::into) + } else { + deserializer.deserialize_map(ProjectedAtifTrajectoryVisitor { scan: self.scan }) + } } } @@ -181,13 +281,13 @@ impl<'de> Visitor<'de> for ProjectedAtifTrajectoryVisitor<'_> { agent = Some(map.next_value::()?); } ProjectedAtifTrajectoryField::Steps => { + let flags = ProjectedAtifScanFlags::new(self.scan); let known_session = session_id.as_deref().filter(|value| !value.is_empty()); if known_session.is_some_and(|value| !self.scan.matches_document(value)) { skipped_steps = map.next_value_seed(CountSequenceSeed)?; steps = Some(Vec::new()); } else { - steps = - Some(map.next_value_seed(ProjectedAtifStepsSeed { scan: self.scan })?); + steps = Some(map.next_value_seed(ProjectedAtifStepsSeed { flags })?); } } ProjectedAtifTrajectoryField::SubagentTrajectories => { @@ -289,26 +389,26 @@ impl<'de> Visitor<'de> for ProjectedAtifTrajectoriesVisitor<'_> { } } -struct ProjectedAtifStepsSeed<'a> { - scan: &'a FileScanSpec, +struct ProjectedAtifStepsSeed { + flags: ProjectedAtifScanFlags, } -impl<'de> DeserializeSeed<'de> for ProjectedAtifStepsSeed<'_> { +impl<'de> DeserializeSeed<'de> for ProjectedAtifStepsSeed { type Value = Vec; fn deserialize(self, deserializer: D) -> std::result::Result where D: serde::Deserializer<'de>, { - deserializer.deserialize_seq(ProjectedAtifStepsVisitor { scan: self.scan }) + deserializer.deserialize_seq(ProjectedAtifStepsVisitor { flags: self.flags }) } } -struct ProjectedAtifStepsVisitor<'a> { - scan: &'a FileScanSpec, +struct ProjectedAtifStepsVisitor { + flags: ProjectedAtifScanFlags, } -impl<'de> Visitor<'de> for ProjectedAtifStepsVisitor<'_> { +impl<'de> Visitor<'de> for ProjectedAtifStepsVisitor { type Value = Vec; fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { @@ -321,7 +421,7 @@ impl<'de> Visitor<'de> for ProjectedAtifStepsVisitor<'_> { { let mut steps = Vec::with_capacity(sequence.size_hint().unwrap_or_default().min(8192)); while let Some(step) = - sequence.next_element_seed(ProjectedAtifStepSeed { scan: self.scan })? + sequence.next_element_seed(ProjectedAtifStepSeed { flags: self.flags })? { steps.push(step); } @@ -329,26 +429,26 @@ impl<'de> Visitor<'de> for ProjectedAtifStepsVisitor<'_> { } } -struct ProjectedAtifStepSeed<'a> { - scan: &'a FileScanSpec, +struct ProjectedAtifStepSeed { + flags: ProjectedAtifScanFlags, } -impl<'de> DeserializeSeed<'de> for ProjectedAtifStepSeed<'_> { +impl<'de> DeserializeSeed<'de> for ProjectedAtifStepSeed { type Value = ProjectedAtifStep; fn deserialize(self, deserializer: D) -> std::result::Result where D: serde::Deserializer<'de>, { - deserializer.deserialize_map(ProjectedAtifStepVisitor { scan: self.scan }) + deserializer.deserialize_map(ProjectedAtifStepVisitor { flags: self.flags }) } } -struct ProjectedAtifStepVisitor<'a> { - scan: &'a FileScanSpec, +struct ProjectedAtifStepVisitor { + flags: ProjectedAtifScanFlags, } -impl<'de> Visitor<'de> for ProjectedAtifStepVisitor<'_> { +impl<'de> Visitor<'de> for ProjectedAtifStepVisitor { type Value = ProjectedAtifStep; fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { @@ -359,30 +459,26 @@ impl<'de> Visitor<'de> for ProjectedAtifStepVisitor<'_> { where A: MapAccess<'de>, { - let mut seen = HashSet::new(); let mut step_id = None; let mut timestamp = None; let mut source = None; let mut model_name = None; - let mut reasoning_effort = None; - let mut message = serde_json::Value::Null; + let mut reasoning_effort_json = None; + let mut message_json = None; let mut message_seen = false; let mut reasoning_content = None; let mut tool_calls_nonempty = false; let mut observation_present = false; - let mut metrics = None; - let mut extra = None; + let mut metrics_json = None; + let mut extra_json = None; let mut llm_call_count = None; let mut is_copied_context = None; - while let Some(field) = map.next_key::()? { - if field != ProjectedAtifStepField::Other && !seen.insert(field) { - return Err(de::Error::duplicate_field(field.name())); - } - match field { + while let Some(key) = map.next_key::()? { + match key { ProjectedAtifStepField::StepId => step_id = Some(map.next_value::()?), ProjectedAtifStepField::Timestamp => { - if self.scan.wants("timestamp") { + if self.flags.timestamp { timestamp = map.next_value::>()?; } else { map.next_value::()?; @@ -390,76 +486,76 @@ impl<'de> Visitor<'de> for ProjectedAtifStepVisitor<'_> { } ProjectedAtifStepField::Source => source = Some(map.next_value::()?), ProjectedAtifStepField::ModelName => { - if self.scan.wants("model_name") { + if self.flags.model_name { model_name = map.next_value::>()?; } else { map.next_value::()?; } } ProjectedAtifStepField::ReasoningEffort => { - if self.scan.wants("reasoning_effort_json") { - reasoning_effort = map.next_value::>()?; + if self.flags.reasoning_effort_json { + reasoning_effort_json = + map.next_value::>>()?; } else { map.next_value::()?; } } ProjectedAtifStepField::Message => { message_seen = true; - if self.scan.wants("message_json") { - message = map.next_value::()?; + if self.flags.message_json { + message_json = Some(map.next_value::>()?); } else { map.next_value::()?; } } ProjectedAtifStepField::ReasoningContent => { - if self.scan.wants("reasoning_content") { + if self.flags.reasoning_content { reasoning_content = map.next_value::>()?; } else { map.next_value::()?; } } ProjectedAtifStepField::ToolCalls => { - if self.scan.wants("kind") || self.scan.wants("effective_kind") { - tool_calls_nonempty = map - .next_value::>>()? - .is_some_and(|calls| !calls.is_empty()); + if self.flags.kind_fields { + let raw = map.next_value::>>()?; + tool_calls_nonempty = + raw.is_some_and(|value| raw_json_value_present(value.get())); } else { map.next_value::()?; } } ProjectedAtifStepField::Observation => { - if self.scan.wants("had_observation") { + if self.flags.had_observation { observation_present = map.next_value::>()?.is_some(); } else { map.next_value::()?; } } ProjectedAtifStepField::Metrics => { - if self.scan.wants("metrics_json") - || self.scan.wants("latency_ms") - || self.scan.wants("ttft_ms") - { - metrics = map.next_value::>()?; + if self.flags.metrics { + metrics_json = + map.next_value::>>()?; } else { map.next_value::()?; } } ProjectedAtifStepField::Extra => { - if self.scan.wants("extra_json") { - extra = map.next_value::>()?; + if self.flags.extra_json { + extra_json = + map.next_value::>>()?; } else { map.next_value::()?; } } ProjectedAtifStepField::LlmCallCount => { - if self.scan.wants("llm_call_count") { + if self.flags.llm_call_count { llm_call_count = map.next_value::>()?; } else { map.next_value::()?; } } ProjectedAtifStepField::IsCopiedContext => { - if self.scan.wants("is_copied_context") { + if self.flags.is_copied_context { is_copied_context = map.next_value::>()?; } else { map.next_value::()?; @@ -478,13 +574,13 @@ impl<'de> Visitor<'de> for ProjectedAtifStepVisitor<'_> { timestamp, source: source.ok_or_else(|| de::Error::missing_field("source"))?, model_name, - reasoning_effort, - message, + reasoning_effort_json, + message_json, reasoning_content, tool_calls_nonempty, observation_present, - metrics, - extra, + metrics_json, + extra_json, llm_call_count, is_copied_context, }) @@ -518,7 +614,10 @@ impl<'de> Visitor<'de> for CountSequenceVisitor { A: SeqAccess<'de>, { let mut count = 0; - while sequence.next_element::()?.is_some() { + while sequence + .next_element::>()? + .is_some() + { count += 1; } Ok(count) @@ -534,7 +633,7 @@ struct ProjectedAtifStream<'a> { batch_size: usize, scan: &'a FileScanSpec, tx: &'a Sender>, - pending: Vec, + pending: Vec, document_ids: HashSet, cancelled: bool, } @@ -600,256 +699,29 @@ impl<'a> ProjectedAtifStream<'a> { } } -struct BoundedCountingReader { - inner: R, - bytes_read: u64, - maximum: u64, -} - -impl BoundedCountingReader { - fn new(inner: R, maximum: u64) -> Self { - Self { - inner, - bytes_read: 0, - maximum, - } - } - - fn bytes_read(&self) -> u64 { - self.bytes_read - } -} - -impl Read for BoundedCountingReader { - fn read(&mut self, buffer: &mut [u8]) -> io::Result { - let remaining = self.maximum.saturating_sub(self.bytes_read); - if remaining == 0 { - let mut probe = [0_u8; 1]; - if self.inner.read(&mut probe)? == 0 { - return Ok(0); - } - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("trajectory input exceeded {} bytes", self.maximum), - )); - } - let maximum = usize::try_from(remaining.min(buffer.len() as u64)).unwrap_or(buffer.len()); - let read = self.inner.read(&mut buffer[..maximum])?; - self.bytes_read += read as u64; - Ok(read) - } -} - -fn read_bounded_line( - reader: &mut R, - buffer: &mut Vec, - maximum: usize, -) -> io::Result { - buffer.clear(); - loop { - let available = reader.fill_buf()?; - if available.is_empty() { - return Ok(buffer.len()); - } - let end = available - .iter() - .position(|byte| *byte == b'\n') - .map_or(available.len(), |index| index + 1); - if buffer.len().saturating_add(end) > maximum { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("JSONL record exceeded max_record_bytes {maximum}"), - )); - } - buffer.extend_from_slice(&available[..end]); - let ended = available[end - 1] == b'\n'; - reader.consume(end); - if ended { - return Ok(buffer.len()); - } - } -} - -/// Copy one complete top-level JSON object out of a buffered stream. -/// -/// This scanner only discovers the record boundary; serde remains the source -/// of truth for JSON syntax and ATIF validation. Strings and escapes are -/// tracked so braces inside message text do not terminate the record. -fn read_bounded_json_object( - reader: &mut R, - buffer: &mut Vec, - maximum: usize, -) -> io::Result { - buffer.clear(); - let mut depth = 0_usize; - let mut in_string = false; - let mut escaped = false; - - loop { - let available = reader.fill_buf()?; - if available.is_empty() { - return Err(io::Error::new( - io::ErrorKind::UnexpectedEof, - "unterminated JSON object in array", - )); - } - let mut end = available.len(); - let mut finished = false; - for (index, byte) in available.iter().copied().enumerate() { - if in_string { - if escaped { - escaped = false; - } else if byte == b'\\' { - escaped = true; - } else if byte == b'"' { - in_string = false; - } - continue; - } - match byte { - b'"' => in_string = true, - b'{' | b'[' => { - depth = depth.checked_add(1).ok_or_else(|| { - io::Error::new(io::ErrorKind::InvalidData, "JSON nesting depth overflow") - })?; - } - b'}' | b']' => { - if depth == 0 { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - "unexpected JSON closing delimiter", - )); - } - depth -= 1; - if depth == 0 { - if byte != b'}' { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - "ATIF array element must be a JSON object", - )); - } - end = index + 1; - finished = true; - break; - } - } - _ => {} - } - } - if buffer.len().saturating_add(end) > maximum { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("JSON array record exceeded max_record_bytes {maximum}"), - )); - } - buffer.extend_from_slice(&available[..end]); - reader.consume(end); - if finished { - return Ok(buffer.len()); - } - } -} - -fn trim_ascii_whitespace(mut input: &[u8]) -> &[u8] { - while input.first().is_some_and(u8::is_ascii_whitespace) { - input = &input[1..]; - } - while input.last().is_some_and(u8::is_ascii_whitespace) { - input = &input[..input.len() - 1]; - } - input -} - -fn first_non_whitespace(reader: &mut R) -> io::Result> { - loop { - let available = reader.fill_buf()?; - if available.is_empty() { - return Ok(None); - } - if let Some(index) = available - .iter() - .position(|byte| !byte.is_ascii_whitespace()) - { - let first = available[index]; - reader.consume(index); - return Ok(Some(first)); - } - let length = available.len(); - reader.consume(length); - } -} - -fn is_ndjson(path: &Path) -> bool { - path.extension() - .and_then(|value| value.to_str()) - .is_some_and(|value| matches!(value.to_ascii_lowercase().as_str(), "jsonl" | "ndjson")) -} - -fn stream_projected_atif_array( +fn consume_projected_atif_reader( reader: &mut R, - reader_capacity: usize, + scan: &FileScanSpec, stream: &mut ProjectedAtifStream<'_>, - maximum_record_bytes: usize, ) -> Result<()> { - anyhow::ensure!( - first_non_whitespace(reader)? == Some(b'['), - "projected ATIF array must start with '['" - ); - reader.consume(1); - - let mut first = true; - let mut ordinal = 0_usize; - let mut record = Vec::new(); - loop { - if !first { - match first_non_whitespace(reader)? { - Some(b']') => { - reader.consume(1); - anyhow::ensure!( - first_non_whitespace(reader)?.is_none(), - "trailing content after ATIF JSON array" - ); - return Ok(()); - } - Some(b',') => reader.consume(1), - Some(other) => { - anyhow::bail!("ATIF JSON array expected ',' or ']', found byte 0x{other:02x}") - } - None => anyhow::bail!("unterminated ATIF JSON array"), - } - } - - match first_non_whitespace(reader)? { - Some(b']') if first => anyhow::bail!("ATIF input contains no trajectories"), - Some(b'{') => {} - Some(other) => { - anyhow::bail!("ATIF JSON array element must be an object, found byte 0x{other:02x}") - } - None => anyhow::bail!("unterminated ATIF JSON array"), - } + let mut deserializer = serde_json::Deserializer::from_reader(reader); + let trajectory = ProjectedAtifTrajectorySeed { scan } + .deserialize(&mut deserializer) + .map_err(anyhow::Error::from)?; + deserializer.end().map_err(anyhow::Error::from)?; + stream.consume(trajectory) +} - ordinal += 1; - read_bounded_json_object(reader, &mut record, maximum_record_bytes) - .with_context(|| format!("read projected ATIF array element {ordinal}"))?; - stream - .runtime - .metrics - .inner - .streaming_buffer_peak_bytes - .fetch_max( - reader_capacity.saturating_add(record.capacity()) as u64, - Ordering::Relaxed, - ); - let mut deserializer = serde_json::Deserializer::from_slice(&record); - let trajectory = ProjectedAtifTrajectorySeed { scan: stream.scan } - .deserialize(&mut deserializer) - .with_context(|| format!("parse projected ATIF array element {ordinal}"))?; - deserializer - .end() - .with_context(|| format!("finish projected ATIF array element {ordinal}"))?; - stream.consume(trajectory)?; - first = false; - } +fn deserialize_projected_atif_from_slice( + record: &[u8], + scan: &FileScanSpec, +) -> Result { + let mut deserializer = serde_json::Deserializer::from_slice(record); + let trajectory = ProjectedAtifTrajectorySeed { scan } + .deserialize(&mut deserializer) + .map_err(anyhow::Error::from)?; + deserializer.end().map_err(anyhow::Error::from)?; + Ok(trajectory) } pub(super) fn stream_projected_atif_steps( @@ -882,92 +754,54 @@ pub(super) fn stream_projected_atif_steps( .fetch_max(reader.capacity() as u64, Ordering::Relaxed); let mut stream = ProjectedAtifStream::new(file, runtime, schema, batch_size, scan, tx); - if is_ndjson(file.file.path()) { - let mut record = Vec::new(); - let mut line_number = 0_usize; - let mut parsed_records = 0_usize; - loop { - let read = - read_bounded_line(&mut reader, &mut record, runtime.options.max_record_bytes) - .with_context(|| { - format!("read projected ATIF JSONL {}", file.file.path().display()) - })?; - if read == 0 { - break; - } - line_number += 1; + let reader_capacity = reader.capacity() as u64; + let max_record_bytes = runtime.options.max_record_bytes; + let result = visit_json_stream( + file.file.path(), + &mut reader, + max_record_bytes, + &mut stream, + // A single object spans the whole file and is already bounded by + // `max_file_bytes`, so deserialize it without an intermediate copy. + |reader, stream| { + consume_projected_atif_reader(reader, scan, stream) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error)) + }, + |record, location, stream| { runtime.metrics.inner.streaming_buffer_peak_bytes.fetch_max( - reader.capacity().saturating_add(record.capacity()) as u64, + reader_capacity.saturating_add(record.len() as u64), Ordering::Relaxed, ); - let record = trim_ascii_whitespace(&record); - if record.is_empty() { - continue; - } - let mut deserializer = serde_json::Deserializer::from_slice(record); - let trajectory = ProjectedAtifTrajectorySeed { scan } - .deserialize(&mut deserializer) + let trajectory = deserialize_projected_atif_from_slice(record, scan) .with_context(|| { format!( - "parse projected ATIF JSONL {} line {line_number}", + "parse projected ATIF {location} in {}", file.file.path().display() ) - })?; - deserializer.end().with_context(|| { - format!( - "finish projected ATIF JSONL {} line {line_number}", - file.file.path().display() - ) - })?; - if let Err(error) = stream.consume(trajectory) { - if stream.cancelled { - break; - } - return Err(error); - } - parsed_records += 1; + }) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + stream + .consume(trajectory) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + Ok(()) + }, + ) + .map(|_| ()) + .map_err(anyhow::Error::from) + .or_else(|error| { + if stream.cancelled { + Ok(()) + } else if error.to_string().contains("JSON array contains no objects") { + Err(anyhow::anyhow!("ATIF input contains no trajectories")) + } else { + Err(error) } - anyhow::ensure!( - parsed_records > 0 || stream.cancelled, - "ATIF input contains no trajectories: {}", - file.file.path().display() - ); - } else { - let shape = first_non_whitespace(&mut reader) - .with_context(|| format!("inspect ATIF input {}", file.file.path().display()))? - .with_context(|| format!("ATIF input is empty: {}", file.file.path().display()))?; - let result = match shape { - b'{' => { - let mut deserializer = serde_json::Deserializer::from_reader(&mut reader); - let result = ProjectedAtifTrajectorySeed { scan } - .deserialize(&mut deserializer) - .map_err(anyhow::Error::from) - .and_then(|trajectory| stream.consume(trajectory)); - match result { - Ok(()) => deserializer.end().map_err(anyhow::Error::from), - Err(error) => Err(error), - } - } - b'[' => { - let reader_capacity = reader.capacity(); - stream_projected_atif_array( - &mut reader, - reader_capacity, - &mut stream, - runtime.options.max_record_bytes, - ) - } - _ => anyhow::bail!( - "ATIF input {} must contain an object, array, JSONL, or NDJSON", - file.file.path().display() - ), - }; - if let Err(error) = result { - if !stream.cancelled { - return Err(error).with_context(|| { - format!("parse projected ATIF input {}", file.file.path().display()) - }); - } + }); + if let Err(error) = result { + if !stream.cancelled { + return Err(error).with_context(|| { + format!("parse projected ATIF input {}", file.file.path().display()) + }); } } stream.finish()?; @@ -1002,7 +836,7 @@ fn project_atif_trajectory( batch_size: usize, scan: &FileScanSpec, tx: &Sender>, - pending: &mut Vec, + pending: &mut Vec, document_ids: &mut HashSet, inherited_session_id: Option<&str>, embedded: bool, @@ -1121,7 +955,7 @@ fn project_atif_step( session_id: &str, step: ProjectedAtifStep, scan: &FileScanSpec, -) -> StoryStepRow { +) -> ProjectedStepRow { let wants_kind = scan.wants("kind"); let wants_effective_kind = scan.wants("effective_kind"); let effective_kind = match step.source.as_str() { @@ -1139,8 +973,9 @@ fn project_atif_step( Some(effective_kind.to_string()) }; - let (latency_ms, ttft_ms) = projected_timing_from_metrics(step.metrics.as_ref()); - StoryStepRow { + let (latency_ms, ttft_ms) = + projected_timing_from_metrics(step.metrics_json.as_ref().map(|value| value.get())); + ProjectedStepRow { document_id: if scan.wants("document_id") { document_id.to_string() } else { @@ -1165,20 +1000,30 @@ fn project_atif_step( } else { String::new() }, - message: if scan.wants("message_json") { - step.message + message_json: if scan.wants("message_json") { + step.message_json + .as_deref() + .map(canonical_json_text) + .unwrap_or_else(|| "null".to_string()) } else { - serde_json::Value::Null + "null".to_string() }, reasoning_content: scan .wants("reasoning_content") .then_some(step.reasoning_content) .flatten(), - reasoning_effort: scan + reasoning_effort_json: scan .wants("reasoning_effort_json") - .then_some(step.reasoning_effort) + .then_some( + step.reasoning_effort_json + .as_deref() + .map(canonical_json_text), + ) + .flatten(), + metrics_json: scan + .wants("metrics_json") + .then_some(step.metrics_json.as_deref().map(canonical_json_text)) .flatten(), - metrics: scan.wants("metrics_json").then_some(step.metrics).flatten(), model_name: scan .wants("model_name") .then_some(step.model_name) @@ -1194,151 +1039,9 @@ fn project_atif_step( latency_ms: scan.wants("latency_ms").then_some(latency_ms).flatten(), ttft_ms: scan.wants("ttft_ms").then_some(ttft_ms).flatten(), had_observation: scan.wants("had_observation") && step.observation_present, - extra: scan.wants("extra_json").then_some(step.extra).flatten(), - } -} - -fn projected_timing_from_metrics( - metrics: Option<&serde_json::Value>, -) -> (Option, Option) { - let Some(metrics) = metrics else { - return (None, None); - }; - let latency_ms = metrics - .get("latency_ms") - .or_else(|| metrics.get("elapsed_ms")) - .or_else(|| metrics.get("duration_ms")) - .and_then(|value| { - value - .as_i64() - .or_else(|| value.as_f64().map(|value| value as i64)) - }); - let ttft_ms = metrics.get("ttft_ms").and_then(|value| { - value - .as_i64() - .or_else(|| value.as_f64().map(|value| value as i64)) - }); - (latency_ms, ttft_ms) -} - -fn emit_projected_step_batch( - rows: &mut Vec, - file: &Arc, - runtime: &Arc, - schema: &SchemaRef, - tx: &Sender>, -) -> Result { - if rows.is_empty() { - return Ok(true); - } - let batch = projected_step_rows_to_batch(rows, file.file.relative_path(), schema.clone())?; - rows.clear(); - runtime - .metrics - .inner - .projected_arrow_bytes - .fetch_add(batch.get_array_memory_size() as u64, Ordering::Relaxed); - Ok(tx.blocking_send(Ok(batch)).is_ok()) -} - -fn projected_step_rows_to_batch( - rows: &[StoryStepRow], - relative_path: &str, - schema: SchemaRef, -) -> Result { - let mut columns = Vec::::with_capacity(schema.fields().len()); - for field in schema.fields() { - let column: ArrayRef = match field.name().as_str() { - "document_id" => Arc::new(StringArray::from_iter_values( - rows.iter().map(|row| row.document_id.as_str()), - )), - "run_id" => Arc::new(StringArray::from_iter( - rows.iter().map(|row| row.run_id.as_deref()), - )), - "session_id" => Arc::new(StringArray::from_iter_values( - rows.iter().map(|row| row.session_id.as_str()), - )), - "step_id" => Arc::new(Int64Array::from( - rows.iter().map(|row| row.step_id).collect::>(), - )), - "kind" => Arc::new(StringArray::from_iter( - rows.iter().map(|row| row.kind.as_deref()), - )), - "effective_kind" => Arc::new(StringArray::from_iter_values( - rows.iter().map(|row| row.effective_kind.as_str()), - )), - "timestamp" => Arc::new(timestamp_array( - rows.iter().map(|row| row.timestamp.as_deref()), - )?), - "timestamp_rfc3339" => Arc::new(StringArray::from_iter( - rows.iter().map(|row| row.timestamp.as_deref()), - )), - "source" => Arc::new(StringArray::from_iter_values( - rows.iter().map(|row| row.source.as_str()), - )), - "message_json" => Arc::new(StringArray::from_iter_values( - rows.iter() - .map(|row| serde_json::to_string(&row.message)) - .collect::>>()? - .iter() - .map(String::as_str), - )), - "reasoning_content" => Arc::new(StringArray::from_iter( - rows.iter().map(|row| row.reasoning_content.as_deref()), - )), - "reasoning_effort_json" => Arc::new(optional_json_array( - rows.iter().map(|row| row.reasoning_effort.as_ref()), - )?), - "metrics_json" => Arc::new(optional_json_array( - rows.iter().map(|row| row.metrics.as_ref()), - )?), - "model_name" => Arc::new(StringArray::from_iter( - rows.iter().map(|row| row.model_name.as_deref()), - )), - "llm_call_count" => Arc::new(Int64Array::from( - rows.iter() - .map(|row| row.llm_call_count) - .collect::>(), - )), - "is_copied_context" => Arc::new(BooleanArray::from( - rows.iter() - .map(|row| row.is_copied_context) - .collect::>(), - )), - "latency_ms" => Arc::new(Int64Array::from( - rows.iter().map(|row| row.latency_ms).collect::>(), - )), - "ttft_ms" => Arc::new(Int64Array::from( - rows.iter().map(|row| row.ttft_ms).collect::>(), - )), - "had_observation" => Arc::new(BooleanArray::from( - rows.iter() - .map(|row| row.had_observation) - .collect::>(), - )), - "extra_json" => Arc::new(optional_json_array( - rows.iter().map(|row| row.extra.as_ref()), - )?), - SOURCE_FILE_COLUMN => Arc::new(StringArray::from_iter_values(std::iter::repeat_n( - relative_path, - rows.len(), - ))), - name => anyhow::bail!("unsupported projected ATIF steps column '{name}'"), - }; - columns.push(column); + extra_json: scan + .wants("extra_json") + .then_some(step.extra_json.as_deref().map(canonical_json_text)) + .flatten(), } - let options = RecordBatchOptions::new().with_row_count(Some(rows.len())); - RecordBatch::try_new_with_options(schema, columns, &options) - .context("build projected ATIF steps batch") -} - -fn optional_json_array<'a>( - values: impl IntoIterator>, -) -> Result { - Ok(StringArray::from( - values - .into_iter() - .map(|value| value.map(serde_json::to_string).transpose()) - .collect::>>()?, - )) } diff --git a/crates/persisting-pchronicle/src/store/files/json_stream.rs b/crates/persisting-pchronicle/src/store/files/json_stream.rs new file mode 100644 index 00000000..d331a947 --- /dev/null +++ b/crates/persisting-pchronicle/src/store/files/json_stream.rs @@ -0,0 +1,609 @@ +use std::io::{self, BufRead, Read}; +use std::path::Path; + +/// Top-level JSON stream shapes supported by local trajectory datasources. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum JsonStreamShape { + Object, + Array, + Ndjson, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum JsonRecordLocation { + ArrayElement(usize), + NdjsonLine(usize), +} + +impl std::fmt::Display for JsonRecordLocation { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::ArrayElement(ordinal) => write!(formatter, "array element {ordinal}"), + Self::NdjsonLine(line) => write!(formatter, "JSONL line {line}"), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) struct JsonStreamVisit { + pub(super) record_count: usize, + pub(super) peak_record_bytes: usize, +} + +/// Bounded reader that tracks source bytes and rejects reads past `maximum`. +pub(super) struct BoundedCountingReader { + inner: R, + bytes_read: u64, + maximum: u64, +} + +impl BoundedCountingReader { + pub(super) fn new(inner: R, maximum: u64) -> Self { + Self { + inner, + bytes_read: 0, + maximum, + } + } + + pub(super) fn bytes_read(&self) -> u64 { + self.bytes_read + } +} + +impl Read for BoundedCountingReader { + fn read(&mut self, buffer: &mut [u8]) -> io::Result { + let remaining = self.maximum.saturating_sub(self.bytes_read); + if remaining == 0 { + let mut probe = [0_u8; 1]; + if self.inner.read(&mut probe)? == 0 { + return Ok(0); + } + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("trajectory input exceeded {} bytes", self.maximum), + )); + } + let maximum = usize::try_from(remaining.min(buffer.len() as u64)).unwrap_or(buffer.len()); + let read = self.inner.read(&mut buffer[..maximum])?; + self.bytes_read += read as u64; + Ok(read) + } +} + +/// Feeds one top-level JSON object from a buffered stream into serde without +/// first copying the whole object into an intermediate `Vec`. +/// +/// The reader tracks object depth and stops at the matching `}`. Bytes consumed +/// are counted against `maximum`; oversized objects fail closed. +pub(super) struct ScopedJsonObjectReader<'a, R> { + inner: &'a mut R, + maximum: usize, + bytes: usize, + depth: usize, + started: bool, + finished: bool, + in_string: bool, + escaped: bool, +} + +impl<'a, R: BufRead> ScopedJsonObjectReader<'a, R> { + pub(super) fn new(inner: &'a mut R, maximum: usize) -> Self { + Self { + inner, + maximum, + bytes: 0, + depth: 0, + started: false, + finished: false, + in_string: false, + escaped: false, + } + } + + pub(super) fn is_finished(&self) -> bool { + self.finished + } +} + +impl Read for ScopedJsonObjectReader<'_, R> { + fn read(&mut self, buffer: &mut [u8]) -> io::Result { + if self.finished || buffer.is_empty() { + return Ok(0); + } + + let available = self.inner.fill_buf()?; + if available.is_empty() { + if self.started && !self.finished { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "unterminated JSON object in array", + )); + } + return Ok(0); + } + + let mut take = 0_usize; + for &byte in available.iter().take(buffer.len()) { + if self.bytes == self.maximum { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "JSON array record exceeded max_record_bytes {}", + self.maximum + ), + )); + } + + take += 1; + self.bytes += 1; + if self.in_string { + if self.escaped { + self.escaped = false; + } else if byte == b'\\' { + self.escaped = true; + } else if byte == b'"' { + self.in_string = false; + } + continue; + } + + match byte { + b'"' => self.in_string = true, + b'{' | b'[' => { + self.started = true; + self.depth = self.depth.checked_add(1).ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidData, "JSON nesting depth overflow") + })?; + } + b'}' | b']' => { + if !self.started || self.depth == 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "unexpected JSON closing delimiter", + )); + } + self.depth -= 1; + if self.depth == 0 { + if byte != b'}' { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "JSON array element must be a JSON object", + )); + } + self.finished = true; + break; + } + } + _ => { + if !self.started { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "JSON array element must be an object", + )); + } + } + } + } + + buffer[..take].copy_from_slice(&available[..take]); + self.inner.consume(take); + Ok(take) + } +} + +/// Reads one top-level JSON object into `record`, reusing the caller's +/// allocation across elements. +/// +/// serde_json's reader-based deserializer pulls a single byte per `Read::read` +/// call, which makes parsing through `Read` several times slower than parsing a +/// slice. Copying one bounded element into a reused buffer keeps peak memory at +/// `maximum` while letting callers use the much faster slice parser. +pub(super) fn read_bounded_json_object( + reader: &mut R, + record: &mut Vec, + maximum: usize, +) -> io::Result<()> { + record.clear(); + let mut depth = 0_usize; + let mut in_string = false; + let mut escaped = false; + loop { + let available = reader.fill_buf()?; + if available.is_empty() { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "unterminated JSON object in array", + )); + } + + // Never scan more than one byte past the budget so an oversized element + // fails without walking the rest of the stream. + let scan_limit = maximum + .saturating_sub(record.len()) + .saturating_add(1) + .min(available.len()); + let mut take = 0_usize; + let mut complete = false; + for &byte in &available[..scan_limit] { + take += 1; + if in_string { + if escaped { + escaped = false; + } else if byte == b'\\' { + escaped = true; + } else if byte == b'"' { + in_string = false; + } + continue; + } + match byte { + b'"' => in_string = true, + b'{' | b'[' => { + depth = depth.checked_add(1).ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidData, "JSON nesting depth overflow") + })?; + } + b'}' | b']' => { + if depth == 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "unexpected JSON closing delimiter", + )); + } + depth -= 1; + if depth == 0 { + if byte != b'}' { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "JSON array element must be a JSON object", + )); + } + complete = true; + break; + } + } + _ => { + if depth == 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "JSON array element must be an object", + )); + } + } + } + } + + if record.len().saturating_add(take) > maximum { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("JSON array record exceeded max_record_bytes {maximum}"), + )); + } + record.extend_from_slice(&available[..take]); + reader.consume(take); + if complete { + return Ok(()); + } + } +} + +pub(super) fn read_bounded_line( + reader: &mut R, + buffer: &mut Vec, + maximum: usize, +) -> io::Result { + buffer.clear(); + loop { + let available = reader.fill_buf()?; + if available.is_empty() { + return Ok(buffer.len()); + } + let end = available + .iter() + .position(|byte| *byte == b'\n') + .map_or(available.len(), |index| index + 1); + if buffer.len().saturating_add(end) > maximum { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("JSONL record exceeded max_record_bytes {maximum}"), + )); + } + buffer.extend_from_slice(&available[..end]); + let ended = available[end - 1] == b'\n'; + reader.consume(end); + if ended { + return Ok(buffer.len()); + } + } +} + +pub(super) fn trim_ascii_whitespace(mut input: &[u8]) -> &[u8] { + while input.first().is_some_and(u8::is_ascii_whitespace) { + input = &input[1..]; + } + while input.last().is_some_and(u8::is_ascii_whitespace) { + input = &input[..input.len() - 1]; + } + input +} + +pub(super) fn first_non_whitespace(reader: &mut R) -> io::Result> { + loop { + let available = reader.fill_buf()?; + if available.is_empty() { + return Ok(None); + } + if let Some(index) = available + .iter() + .position(|byte| !byte.is_ascii_whitespace()) + { + let first = available[index]; + reader.consume(index); + return Ok(Some(first)); + } + let length = available.len(); + reader.consume(length); + } +} + +pub(super) fn is_ndjson(path: &Path) -> bool { + path.extension() + .and_then(|value| value.to_str()) + .is_some_and(|value| matches!(value.to_ascii_lowercase().as_str(), "jsonl" | "ndjson")) +} + +pub(super) fn detect_json_stream_shape( + path: &Path, + reader: &mut R, +) -> io::Result { + if is_ndjson(path) { + return Ok(JsonStreamShape::Ndjson); + } + match first_non_whitespace(reader)? { + Some(b'{') => Ok(JsonStreamShape::Object), + Some(b'[') => Ok(JsonStreamShape::Array), + Some(other) => Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("unsupported JSON stream start byte 0x{other:02x}"), + )), + None => Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "JSON input is empty", + )), + } +} + +/// Dispatch all supported top-level JSON shapes while keeping array and +/// NDJSON records bounded. The object callback receives the original reader, +/// allowing callers to deserialize it directly without an intermediate copy. +pub(super) fn visit_json_stream( + path: &Path, + reader: &mut R, + max_record_bytes: usize, + state: &mut S, + visit_object: O, + mut visit_record: F, +) -> io::Result +where + R: BufRead, + O: FnOnce(&mut R, &mut S) -> io::Result<()>, + F: FnMut(&[u8], JsonRecordLocation, &mut S) -> io::Result<()>, +{ + let shape = detect_json_stream_shape(path, reader)?; + let mut peak_record_bytes = 0_usize; + let record_count = match shape { + JsonStreamShape::Object => { + visit_object(reader, state)?; + 1 + } + JsonStreamShape::Array => { + for_each_json_array_record(reader, max_record_bytes, |record, ordinal| { + peak_record_bytes = peak_record_bytes.max(record.len()); + visit_record(record, JsonRecordLocation::ArrayElement(ordinal), state) + })? + } + JsonStreamShape::Ndjson => { + for_each_ndjson_line(reader, max_record_bytes, |record, line| { + peak_record_bytes = peak_record_bytes.max(record.len()); + visit_record(record, JsonRecordLocation::NdjsonLine(line), state) + })? + } + }; + Ok(JsonStreamVisit { + record_count, + peak_record_bytes, + }) +} + +/// Stream each top-level object in a JSON array through `visit`, one bounded +/// element at a time. +/// +/// `visit` receives the raw bytes of a single element, valid until the next +/// iteration. Peak memory stays within `max_record_bytes` because the backing +/// buffer is reused across elements. +pub(super) fn for_each_json_array_record( + reader: &mut R, + max_record_bytes: usize, + mut visit: F, +) -> io::Result +where + R: BufRead, + F: FnMut(&[u8], usize) -> io::Result<()>, +{ + let mut record = Vec::new(); + anyhow_io_ensure( + first_non_whitespace(reader)? == Some(b'['), + "JSON array must start with '['", + )?; + reader.consume(1); + + let mut first = true; + let mut ordinal = 0_usize; + loop { + if !first { + match first_non_whitespace(reader)? { + Some(b']') => { + reader.consume(1); + anyhow_io_ensure( + first_non_whitespace(reader)?.is_none(), + "trailing content after JSON array", + )?; + return Ok(ordinal); + } + Some(b',') => reader.consume(1), + Some(other) => { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("JSON array expected ',' or ']', found byte 0x{other:02x}"), + )); + } + None => { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "unterminated JSON array", + )); + } + } + } + + match first_non_whitespace(reader)? { + Some(b']') if first => { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "JSON array contains no objects", + )); + } + Some(b'{') => {} + Some(other) => { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("JSON array element must be an object, found byte 0x{other:02x}"), + )); + } + None => { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "unterminated JSON array", + )); + } + } + + ordinal += 1; + read_bounded_json_object(reader, &mut record, max_record_bytes)?; + visit(&record, ordinal)?; + first = false; + } +} + +/// Stream non-empty NDJSON/JSONL records through `visit`. +pub(super) fn for_each_ndjson_line( + reader: &mut R, + max_record_bytes: usize, + mut visit: F, +) -> io::Result +where + R: BufRead, + F: FnMut(&[u8], usize) -> io::Result<()>, +{ + let mut buffer = Vec::new(); + let mut line_number = 0_usize; + let mut count = 0_usize; + loop { + let read = read_bounded_line(reader, &mut buffer, max_record_bytes)?; + if read == 0 { + break; + } + line_number += 1; + let record = trim_ascii_whitespace(&buffer); + if record.is_empty() { + continue; + } + count += 1; + visit(record, line_number)?; + } + if count == 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "JSON input contains no objects", + )); + } + Ok(count) +} + +fn anyhow_io_ensure(condition: bool, message: &str) -> io::Result<()> { + if condition { + Ok(()) + } else { + Err(io::Error::new( + io::ErrorKind::InvalidData, + message.to_string(), + )) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Cursor; + + #[test] + fn scoped_object_reader_stops_at_object_end_without_extra_copy() { + let mut input = Cursor::new(br#"{"a":1},{"b":2}]"#); + let mut scoped = ScopedJsonObjectReader::new(&mut input, 64); + let mut buffer = Vec::new(); + scoped.read_to_end(&mut buffer).unwrap(); + assert_eq!(buffer, br#"{"a":1}"#); + assert!(scoped.finished); + assert_eq!(input.position(), 7); + } + + #[test] + fn scoped_object_reader_enforces_max_record_bytes() { + let mut input = Cursor::new(br#"{"message":"hello-world"}]"#); + let mut scoped = ScopedJsonObjectReader::new(&mut input, 8); + let error = scoped.read_to_end(&mut Vec::new()).unwrap_err(); + assert!(error.to_string().contains("max_record_bytes 8")); + } + + #[test] + fn for_each_json_array_record_yields_one_element_at_a_time() { + let mut input = Cursor::new(br#"[ {"a":1}, {"b":2} ]"#); + let mut values = Vec::new(); + let count = for_each_json_array_record(&mut input, 64, |record, _ordinal| { + values.push(serde_json::from_slice::(record).unwrap()); + Ok(()) + }) + .unwrap(); + assert_eq!(count, 2); + assert_eq!(values[0]["a"], 1); + assert_eq!(values[1]["b"], 2); + } + + #[test] + fn read_bounded_json_object_reuses_the_record_buffer() { + let mut input = Cursor::new(br#"{"a":{"nested":[1,2]}}{"b":2}"#); + let mut record = Vec::new(); + read_bounded_json_object(&mut input, &mut record, 64).unwrap(); + assert_eq!(record, br#"{"a":{"nested":[1,2]}}"#); + let capacity = record.capacity(); + read_bounded_json_object(&mut input, &mut record, 64).unwrap(); + assert_eq!(record, br#"{"b":2}"#); + assert_eq!(record.capacity(), capacity); + } + + #[test] + fn read_bounded_json_object_enforces_max_record_bytes() { + let mut input = Cursor::new(br#"{"message":"hello-world"}]"#); + let error = read_bounded_json_object(&mut input, &mut Vec::new(), 8).unwrap_err(); + assert!(error.to_string().contains("max_record_bytes 8")); + } + + #[test] + fn read_bounded_json_object_keeps_braces_inside_strings_out_of_depth() { + let mut input = Cursor::new(br#"{"text":"a}b{c\"}"}"#); + let mut record = Vec::new(); + read_bounded_json_object(&mut input, &mut record, 64).unwrap(); + assert_eq!(record, br#"{"text":"a}b{c\"}"}"#); + } +} diff --git a/crates/persisting-pchronicle/src/store/files/mod.rs b/crates/persisting-pchronicle/src/store/files/mod.rs index de8c9eae..da3e8b37 100644 --- a/crates/persisting-pchronicle/src/store/files/mod.rs +++ b/crates/persisting-pchronicle/src/store/files/mod.rs @@ -3,17 +3,24 @@ //! Each source file is one streaming partition. Query-only `_file_` predicates //! are evaluated against the frozen manifest before partitions are opened. +mod actf_reader; +mod actf_stream; mod atif_reader; mod atif_stream; +mod json_stream; +mod projected_steps; -pub(crate) use atif_reader::parse_storylines as parse_atif_storylines; +use actf_reader::parse_actf_storylines_from_reader_with_stats; +use actf_stream::stream_projected_actf_steps; +use atif_reader::parse_atif_storylines_from_reader_with_stats; pub(crate) use atif_reader::AtifReader; use atif_stream::stream_projected_atif_steps; +use json_stream::BoundedCountingReader; use std::collections::{HashMap, HashSet, VecDeque}; use std::fmt; use std::fs::{self, File}; -use std::io::{self, BufRead, BufReader, Read}; +use std::io::BufReader; use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Condvar, Mutex, MutexGuard, Weak}; @@ -30,19 +37,14 @@ use datafusion::physical_plan::stream::RecordBatchReceiverStreamBuilder; use datafusion::physical_plan::streaming::PartitionStream; use datafusion::physical_plan::{ExecutionPlan, SendableRecordBatchStream}; use datafusion::prelude::SessionContext; -use lance::deps::arrow_array::{ - ArrayRef, BooleanArray, Int64Array, RecordBatch, RecordBatchOptions, StringArray, -}; +use lance::deps::arrow_array::{RecordBatch, StringArray}; use lance::deps::arrow_schema::{DataType, Field, Schema as ArrowSchema, SchemaRef}; use serde::de::{self, DeserializeSeed, IgnoredAny, MapAccess, SeqAccess, Visitor}; use tokio::sync::mpsc::Sender; -use crate::convert::actf_to_storylines; use crate::format::DocumentFormat; -use crate::formats::actf::ActfDocument; use crate::formats::parse_openai_msg_corpus_value; -use super::storyline::rows::timestamp_array; use super::{ datafusion_bridge::{from_datafusion, into_datafusion}, split_storyline, story_runs_arrow_schema, story_runs_to_batch, story_steps_arrow_schema, @@ -303,7 +305,7 @@ impl FileScanSpec { } } - fn can_project_atif_steps(&self, schema: &SchemaRef) -> bool { + fn can_project_steps(&self, schema: &SchemaRef) -> bool { self.projection .as_ref() .is_some_and(|projection| projection.len() < schema.fields().len()) @@ -483,13 +485,13 @@ impl TableProvider for FileTrajectoryTableProvider { .map(|filter| { if matches_file_filter(filter, "").is_some() { TableProviderFilterPushDown::Exact - } else if self.format == DocumentFormat::Atif + } else if matches!(self.format, DocumentFormat::Atif | DocumentFormat::Actf) && self.kind == StorylineTableKind::Steps && atif_step_filters(filter).is_some() { // The projected decoder applies these filters to reduce // materialization, while DataFusion retains the filter to - // guarantee SQL semantics on the compatibility fallback. + // guarantee SQL semantics on the full-normalization path. TableProviderFilterPushDown::Inexact } else { TableProviderFilterPushDown::Unsupported @@ -559,11 +561,18 @@ fn stream_file( ) -> Result<()> { if format == DocumentFormat::Atif && kind == StorylineTableKind::Steps - && scan.can_project_atif_steps(&source_schema) + && scan.can_project_steps(&source_schema) { stream_projected_atif_steps(file, runtime, &schema, batch_size, scan, tx)?; return Ok(()); } + if format == DocumentFormat::Actf + && kind == StorylineTableKind::Steps + && scan.can_project_steps(&source_schema) + { + stream_projected_actf_steps(file, runtime, &schema, batch_size, scan, tx)?; + return Ok(()); + } let parsed = load_file(file, runtime, format)?; for batch in parsed.batches(kind) { let batch = match &scan.projection { @@ -789,32 +798,74 @@ fn load_file( state.file.size_bytes(), runtime.options.max_file_bytes ); - let content = fs::read_to_string(state.file.path()).with_context(|| { - format!( - "read {} input {}", - format.as_str(), - state.file.path().display() - ) - })?; - anyhow::ensure!( - content.len() as u64 <= runtime.options.max_file_bytes, - "{} input {} exceeded max_file_bytes {} while reading", - format.as_str(), - state.file.path().display(), - runtime.options.max_file_bytes - ); - state.file.validate_unchanged()?; - runtime - .metrics - .inner - .source_bytes_read - .fetch_add(content.len() as u64, Ordering::Relaxed); - let parsed = Arc::new(parse_file( - &state.file, - format, - &content, - runtime.options.batch_size, - )?); + let parsed = match format { + DocumentFormat::Atif | DocumentFormat::Actf => { + let input = File::open(state.file.path()).with_context(|| { + format!( + "open {} input {}", + format.as_str(), + state.file.path().display() + ) + })?; + let mut reader = BufReader::with_capacity( + 64 * 1024, + BoundedCountingReader::new(input, runtime.options.max_file_bytes), + ); + let (stories, peak_record_bytes) = match format { + DocumentFormat::Atif => parse_atif_storylines_from_reader_with_stats( + state.file.path(), + &mut reader, + runtime.options.max_record_bytes, + ) + .with_context(|| format!("parse ATIF input {}", state.file.path().display()))?, + DocumentFormat::Actf => parse_actf_storylines_from_reader_with_stats( + state.file.path(), + &mut reader, + runtime.options.max_record_bytes, + ) + .with_context(|| format!("parse ACTF input {}", state.file.path().display()))?, + _ => unreachable!(), + }; + runtime.metrics.inner.streaming_buffer_peak_bytes.fetch_max( + (reader.capacity() as u64).saturating_add(peak_record_bytes as u64), + Ordering::Relaxed, + ); + state.file.validate_unchanged()?; + runtime + .metrics + .inner + .source_bytes_read + .fetch_add(reader.get_ref().bytes_read(), Ordering::Relaxed); + stories_to_parsed_file(&state.file, format, stories, runtime.options.batch_size)? + } + DocumentFormat::OpenaiMsg => { + let content = fs::read_to_string(state.file.path()).with_context(|| { + format!( + "read {} input {}", + format.as_str(), + state.file.path().display() + ) + })?; + anyhow::ensure!( + content.len() as u64 <= runtime.options.max_file_bytes, + "{} input {} exceeded max_file_bytes {} while reading", + format.as_str(), + state.file.path().display(), + runtime.options.max_file_bytes + ); + state.file.validate_unchanged()?; + runtime + .metrics + .inner + .source_bytes_read + .fetch_add(content.len() as u64, Ordering::Relaxed); + parse_openai_stories_from_content(&state.file, &content, runtime.options.batch_size)? + } + unsupported => { + anyhow::bail!("file trajectory datasource does not support '{unsupported}'") + } + }; + let parsed = Arc::new(parsed); runtime .metrics .inner @@ -834,34 +885,25 @@ fn load_file( Ok(parsed) } -fn parse_file( +fn parse_openai_stories_from_content( file: &LocalQueryInputFile, - format: DocumentFormat, content: &str, batch_size: usize, ) -> Result { - let stories = match format { - DocumentFormat::Atif => parse_atif_storylines(content) - .with_context(|| format!("parse ATIF input {}", file.path().display()))?, - DocumentFormat::OpenaiMsg => { - let value = serde_json::from_str(content) - .with_context(|| format!("parse OpenAI JSON input {}", file.path().display()))?; - parse_openai_msg_corpus_value(&value, file.relative_path()) - .map_err(anyhow::Error::from) - .with_context(|| format!("normalize OpenAI input {}", file.path().display()))? - } - DocumentFormat::Actf => { - let document = ActfDocument::from_json_str(content) - .map_err(anyhow::Error::from) - .with_context(|| format!("parse ACTF input {}", file.path().display()))?; - actf_to_storylines(&document) - .with_context(|| format!("normalize ACTF input {}", file.path().display()))? - } - unsupported => { - anyhow::bail!("file trajectory datasource does not support '{unsupported}'") - } - }; + let value = serde_json::from_str(content) + .with_context(|| format!("parse OpenAI JSON input {}", file.path().display()))?; + let stories = parse_openai_msg_corpus_value(&value, file.relative_path()) + .map_err(anyhow::Error::from) + .with_context(|| format!("normalize OpenAI input {}", file.path().display()))?; + stories_to_parsed_file(file, DocumentFormat::OpenaiMsg, stories, batch_size) +} +fn stories_to_parsed_file( + file: &LocalQueryInputFile, + format: DocumentFormat, + stories: Vec, + batch_size: usize, +) -> Result { let mut document_ids = HashSet::with_capacity(stories.len()); let mut runs = Vec::::with_capacity(stories.len()); let mut steps = Vec::::new(); diff --git a/crates/persisting-pchronicle/src/store/files/projected_steps.rs b/crates/persisting-pchronicle/src/store/files/projected_steps.rs new file mode 100644 index 00000000..58484723 --- /dev/null +++ b/crates/persisting-pchronicle/src/store/files/projected_steps.rs @@ -0,0 +1,204 @@ +//! Shared projected step row batching for ATIF and ACTF streaming queries. + +use std::sync::atomic::Ordering; +use std::sync::Arc; + +use anyhow::{Context, Result}; +use lance::deps::arrow_array::{ + ArrayRef, BooleanArray, Int64Array, RecordBatch, RecordBatchOptions, StringArray, +}; +use lance::deps::arrow_schema::SchemaRef; +use tokio::sync::mpsc::Sender; + +use super::{FileState, FileTrajectoryRuntime, SOURCE_FILE_COLUMN}; +use crate::store::storyline::rows::timestamp_array; + +pub(crate) struct ProjectedStepRow { + pub document_id: String, + pub run_id: Option, + pub session_id: String, + pub step_id: i64, + pub kind: Option, + pub effective_kind: String, + pub timestamp: Option, + pub source: String, + pub message_json: String, + pub reasoning_content: Option, + pub reasoning_effort_json: Option, + pub metrics_json: Option, + pub model_name: Option, + pub llm_call_count: Option, + pub is_copied_context: Option, + pub latency_ms: Option, + pub ttft_ms: Option, + pub had_observation: bool, + pub extra_json: Option, +} + +/// Renders JSON captured as a [`serde_json::value::RawValue`] in the same +/// canonical form the fully materialized path emits. +/// +/// Raw capture keeps the source bytes verbatim, so without this the emitted +/// column would echo the input file's indentation and key order and identical +/// data stored in differently formatted files would compare unequal. +pub(crate) fn canonical_json_text(raw: &serde_json::value::RawValue) -> String { + match serde_json::from_str::(raw.get()) { + Ok(value) => value.to_string(), + // The text came from a completed parse, so this is unreachable; fall + // back to the raw bytes rather than failing the whole scan. + Err(_) => raw.get().to_string(), + } +} + +pub(crate) fn projected_timing_from_metrics( + metrics_json: Option<&str>, +) -> (Option, Option) { + let Some(metrics_json) = metrics_json else { + return (None, None); + }; + #[derive(serde::Deserialize)] + struct MetricsProjection { + latency_ms: Option, + elapsed_ms: Option, + duration_ms: Option, + ttft_ms: Option, + } + let Ok(metrics) = serde_json::from_str::(metrics_json) else { + return (None, None); + }; + let parse = |value: Option| { + value.and_then(|number| { + number + .as_i64() + .or_else(|| number.as_f64().map(|value| value as i64)) + }) + }; + let latency_ms = parse(metrics.latency_ms) + .or_else(|| parse(metrics.elapsed_ms)) + .or_else(|| parse(metrics.duration_ms)); + let ttft_ms = parse(metrics.ttft_ms); + (latency_ms, ttft_ms) +} + +pub(crate) fn projected_timing_from_actf_metrics(metrics_json: Option<&str>) -> Option { + let metrics_json = metrics_json?; + #[derive(serde::Deserialize)] + struct ActfMetricProjection { + llm_infer_ms: Option, + } + let Ok(metrics) = serde_json::from_str::(metrics_json) else { + return None; + }; + metrics.llm_infer_ms.and_then(|value| { + value + .as_i64() + .or_else(|| value.as_f64().map(|value| value as i64)) + }) +} + +pub(crate) fn emit_projected_step_batch( + rows: &mut Vec, + file: &Arc, + runtime: &Arc, + schema: &SchemaRef, + tx: &Sender>, +) -> Result { + if rows.is_empty() { + return Ok(true); + } + let batch = projected_step_rows_to_batch(rows, file.file.relative_path(), schema.clone())?; + rows.clear(); + runtime + .metrics + .inner + .projected_arrow_bytes + .fetch_add(batch.get_array_memory_size() as u64, Ordering::Relaxed); + Ok(tx.blocking_send(Ok(batch)).is_ok()) +} + +pub(crate) fn projected_step_rows_to_batch( + rows: &[ProjectedStepRow], + relative_path: &str, + schema: SchemaRef, +) -> Result { + let mut columns = Vec::::with_capacity(schema.fields().len()); + for field in schema.fields() { + let column: ArrayRef = match field.name().as_str() { + "document_id" => Arc::new(StringArray::from_iter_values( + rows.iter().map(|row| row.document_id.as_str()), + )), + "session_id" => Arc::new(StringArray::from_iter_values( + rows.iter().map(|row| row.session_id.as_str()), + )), + "run_id" => Arc::new(StringArray::from_iter( + rows.iter().map(|row| row.run_id.as_deref()), + )), + "step_id" => Arc::new(Int64Array::from( + rows.iter().map(|row| row.step_id).collect::>(), + )), + "kind" => Arc::new(StringArray::from_iter( + rows.iter().map(|row| row.kind.as_deref()), + )), + "effective_kind" => Arc::new(StringArray::from_iter_values( + rows.iter().map(|row| row.effective_kind.as_str()), + )), + "timestamp" => Arc::new(timestamp_array( + rows.iter().map(|row| row.timestamp.as_deref()), + )?), + "timestamp_rfc3339" => Arc::new(StringArray::from_iter( + rows.iter().map(|row| row.timestamp.as_deref()), + )), + "source" => Arc::new(StringArray::from_iter_values( + rows.iter().map(|row| row.source.as_str()), + )), + "message_json" => Arc::new(StringArray::from_iter_values( + rows.iter().map(|row| row.message_json.as_str()), + )), + "reasoning_content" => Arc::new(StringArray::from_iter( + rows.iter().map(|row| row.reasoning_content.as_deref()), + )), + "reasoning_effort_json" => Arc::new(StringArray::from_iter( + rows.iter().map(|row| row.reasoning_effort_json.as_deref()), + )), + "metrics_json" => Arc::new(StringArray::from_iter( + rows.iter().map(|row| row.metrics_json.as_deref()), + )), + "model_name" => Arc::new(StringArray::from_iter( + rows.iter().map(|row| row.model_name.as_deref()), + )), + "llm_call_count" => Arc::new(Int64Array::from( + rows.iter() + .map(|row| row.llm_call_count) + .collect::>(), + )), + "is_copied_context" => Arc::new(BooleanArray::from( + rows.iter() + .map(|row| row.is_copied_context) + .collect::>(), + )), + "latency_ms" => Arc::new(Int64Array::from( + rows.iter().map(|row| row.latency_ms).collect::>(), + )), + "ttft_ms" => Arc::new(Int64Array::from( + rows.iter().map(|row| row.ttft_ms).collect::>(), + )), + "had_observation" => Arc::new(BooleanArray::from( + rows.iter() + .map(|row| row.had_observation) + .collect::>(), + )), + "extra_json" => Arc::new(StringArray::from_iter( + rows.iter().map(|row| row.extra_json.as_deref()), + )), + SOURCE_FILE_COLUMN => Arc::new(StringArray::from_iter_values(std::iter::repeat_n( + relative_path, + rows.len(), + ))), + name => anyhow::bail!("unsupported projected steps column '{name}'"), + }; + columns.push(column); + } + let options = RecordBatchOptions::new().with_row_count(Some(rows.len())); + RecordBatch::try_new_with_options(schema, columns, &options) + .context("build projected steps batch") +} diff --git a/crates/persisting-pchronicle/src/store/files/tests.rs b/crates/persisting-pchronicle/src/store/files/tests.rs index f30e081a..c1d3bfe1 100644 --- a/crates/persisting-pchronicle/src/store/files/tests.rs +++ b/crates/persisting-pchronicle/src/store/files/tests.rs @@ -117,6 +117,120 @@ async fn projected_ndjson_enforces_private_record_bound() { ); } +#[tokio::test] +async fn projected_array_enforces_private_record_bound_per_element() { + let trajectory = std::fs::read_to_string(atif_fixture("dialogue_10.json")).unwrap(); + let value: serde_json::Value = serde_json::from_str(&trajectory).unwrap(); + let input = tempfile::NamedTempFile::with_suffix(".json").unwrap(); + std::fs::write( + input.path(), + format!("[{}]", serde_json::to_string(&value).unwrap()), + ) + .unwrap(); + let manifest = LocalQueryManifest::for_format(input.path(), DocumentFormat::Atif).unwrap(); + let source = FileTrajectoryDataSource::from_manifest_with_options( + manifest, + FileTrajectoryDataSourceOptions { + max_record_bytes: 512, + ..FileTrajectoryDataSourceOptions::default() + }, + ) + .unwrap(); + let context = SessionContext::new(); + source.register(&context).unwrap(); + let error = context + .sql("SELECT COUNT(*) FROM steps") + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!( + format!("{error:#}").contains("max_record_bytes 512"), + "{error:#}" + ); +} + +#[tokio::test] +async fn full_atif_array_enforces_private_record_bound_per_element() { + let trajectory = std::fs::read_to_string(atif_fixture("dialogue_10.json")).unwrap(); + let value: serde_json::Value = serde_json::from_str(&trajectory).unwrap(); + let input = tempfile::NamedTempFile::with_suffix(".json").unwrap(); + std::fs::write( + input.path(), + format!("[{}]", serde_json::to_string(&value).unwrap()), + ) + .unwrap(); + let manifest = LocalQueryManifest::for_format(input.path(), DocumentFormat::Atif).unwrap(); + let source = FileTrajectoryDataSource::from_manifest_with_options( + manifest, + FileTrajectoryDataSourceOptions { + max_record_bytes: 512, + ..FileTrajectoryDataSourceOptions::default() + }, + ) + .unwrap(); + let context = SessionContext::new(); + source.register(&context).unwrap(); + let error = context + .sql("SELECT * FROM steps") + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!( + format!("{error:#}").contains("max_record_bytes 512"), + "{error:#}" + ); +} + +#[tokio::test] +async fn full_atif_array_reports_bounded_input_buffer_peak() { + let trajectory = std::fs::read_to_string(atif_fixture("dialogue_10.json")).unwrap(); + let input = tempfile::NamedTempFile::with_suffix(".json").unwrap(); + std::fs::write(input.path(), format!("[{trajectory}]")).unwrap(); + let manifest = LocalQueryManifest::for_format(input.path(), DocumentFormat::Atif).unwrap(); + let source = FileTrajectoryDataSource::from_manifest(manifest).unwrap(); + let context = SessionContext::new(); + source.register(&context).unwrap(); + context + .sql("SELECT * FROM steps") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let peak = source.metrics().snapshot().streaming_buffer_peak_bytes; + assert!(peak >= 64 * 1024, "peak={peak}"); + assert!(peak < 2 * 64 * 1024, "peak={peak}"); +} + +#[tokio::test] +async fn projected_actf_pushdown_matches_session_id_and_step_id() { + let path = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("tests/fixtures/import_roundtrip/protein-assembly_trimmed.actf.json"); + let manifest = LocalQueryManifest::for_format(&path, DocumentFormat::Actf).unwrap(); + let source = FileTrajectoryDataSource::from_manifest(manifest).unwrap(); + let context = SessionContext::new(); + source.register(&context).unwrap(); + + let batches = context + .sql( + "SELECT session_id, step_id, source FROM steps WHERE session_id = 'protein-assembly-trimmed' AND step_id = 2", + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + let metrics = source.metrics().snapshot(); + assert!(metrics.rows_pruned > 0, "{metrics:?}"); +} + #[tokio::test] async fn projected_atif_pushdown_matches_session_id_not_document_id() { let input = tempfile::NamedTempFile::with_suffix(".json").unwrap(); diff --git a/crates/persisting-pchronicle/src/store/storyline/content.rs b/crates/persisting-pchronicle/src/store/storyline/content.rs index 5980c324..881f8d5a 100644 --- a/crates/persisting-pchronicle/src/store/storyline/content.rs +++ b/crates/persisting-pchronicle/src/store/storyline/content.rs @@ -25,6 +25,10 @@ use lance_index::scalar::{BuiltinIndexType, ScalarIndexParams}; use lance_index::IndexType; use super::datafusion::StorylineTableKind; +use crate::formats::unknown_fields::{ + StorylineUnknownFields, UnknownFieldLimits, DEFAULT_MAX_UNKNOWN_BYTES, + DEFAULT_MAX_UNKNOWN_FIELDS, +}; pub const STORYLINE_OBJECTS_DATASET: &str = "objects.lance"; pub const DEFAULT_CONTENT_OFFLOAD_THRESHOLD: usize = 64 * 1024; @@ -54,6 +58,10 @@ pub struct StorylineContentOptions { pub max_chunk_bytes: Option, /// Maximum number of documents accepted by one streamed import. pub max_import_documents: Option, + /// Maximum number of logical unknown fields retained by one document. + pub max_unknown_fields: usize, + /// Maximum logical JSON bytes retained in unknown fields by one document. + pub max_unknown_bytes: usize, } impl Default for StorylineContentOptions { @@ -67,6 +75,8 @@ impl Default for StorylineContentOptions { max_chunk_rows: None, max_chunk_bytes: None, max_import_documents: None, + max_unknown_fields: DEFAULT_MAX_UNKNOWN_FIELDS, + max_unknown_bytes: DEFAULT_MAX_UNKNOWN_BYTES, } } } @@ -92,8 +102,16 @@ impl StorylineContentOptions { anyhow::ensure!(value > 0, "{name} must be positive"); } } + self.unknown_field_limits().validate()?; Ok(self) } + + pub(crate) fn unknown_field_limits(self) -> UnknownFieldLimits { + UnknownFieldLimits { + max_fields: self.max_unknown_fields, + max_bytes: self.max_unknown_bytes, + } + } } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -254,10 +272,34 @@ impl PendingContent { #[derive(Debug, Clone)] struct ResolvedObject { + codec: ContentCodec, raw_length: u64, bytes: Vec, } +pub(crate) fn externalize_unknown_field_values( + fields: &mut StorylineUnknownFields, + options: StorylineContentOptions, + pending: &mut PendingContent, +) -> Result<()> { + for source in fields.sources.values_mut() { + for value in source.fields.values_mut() { + let encoded = serde_json::to_vec(value).context("serialize Storyline unknown value")?; + let collides_with_descriptor = value + .as_str() + .is_some_and(|value| value.starts_with(CONTENT_REF_MAGIC)); + if encoded.len() < options.offload_threshold && !collides_with_descriptor { + continue; + } + let object = build_object(&encoded, LogicalType::Json, options)?; + let descriptor = object.reference.encode(); + pending.insert(object)?; + *value = serde_json::Value::String(descriptor); + } + } + Ok(()) +} + pub(crate) fn content_columns(kind: StorylineTableKind) -> &'static [(&'static str, bool)] { match kind { StorylineTableKind::Runs => &[ @@ -627,6 +669,29 @@ pub(crate) fn collect_content_ids( } } } + if kind == StorylineTableKind::Runs { + let Some(column) = batch.column_by_name("unknown_fields_json") else { + continue; + }; + let values = column + .as_any() + .downcast_ref::() + .context("Storyline unknown_fields_json is not Utf8")?; + for encoded_fields in values.iter().flatten() { + let fields: StorylineUnknownFields = serde_json::from_str(encoded_fields) + .context("decode Storyline unknown_fields_json for content collection")?; + for source in fields.sources.values() { + for value in source.fields.values() { + let Some(encoded) = value.as_str() else { + continue; + }; + if let Some(reference) = ContentRef::parse(encoded)? { + ids.insert(reference.content_id); + } + } + } + } + } } Ok(ids) } @@ -673,7 +738,132 @@ pub(crate) async fn hydrate_batches( .iter() .map(|(name, _)| *name) .collect::>(); - hydrate_selected_batches(dataset, batches, &selected).await + let batches = hydrate_selected_batches(dataset, batches, &selected).await?; + if kind == StorylineTableKind::Runs { + hydrate_unknown_field_values(dataset, batches).await + } else { + Ok(batches) + } +} + +async fn hydrate_unknown_field_values( + dataset: &Arc, + batches: Vec, +) -> Result> { + const COLUMN: &str = "unknown_fields_json"; + let mut references = HashMap::::new(); + for batch in &batches { + let Some(column) = batch.column_by_name(COLUMN) else { + continue; + }; + let values = column + .as_any() + .downcast_ref::() + .context("Storyline unknown_fields_json is not Utf8")?; + for encoded_fields in values.iter().flatten() { + let fields: StorylineUnknownFields = serde_json::from_str(encoded_fields) + .context("decode Storyline unknown_fields_json for hydration")?; + for source in fields.sources.values() { + for value in source.fields.values() { + let Some(encoded) = value.as_str() else { + continue; + }; + let Some(reference) = ContentRef::parse(encoded) + .context("invalid internal unknown-field content descriptor")? + else { + continue; + }; + anyhow::ensure!( + reference.logical_type == LogicalType::Json, + "unknown-field content descriptor is not JSON" + ); + if let Some(existing) = references.get(&reference.content_id) { + anyhow::ensure!( + existing == &reference, + "conflicting Storyline content descriptors for '{}'", + reference.content_id + ); + } else { + references.insert(reference.content_id.clone(), reference); + } + } + } + } + } + if references.is_empty() { + return Ok(batches); + } + let resolved = resolve_objects(dataset, &references).await?; + batches + .into_iter() + .map(|batch| hydrate_unknown_field_batch(batch, &resolved)) + .collect() +} + +fn hydrate_unknown_field_batch( + batch: RecordBatch, + resolved: &HashMap, +) -> Result { + const COLUMN: &str = "unknown_fields_json"; + let Ok(index) = batch.schema().index_of(COLUMN) else { + return Ok(batch); + }; + let values = batch + .column(index) + .as_any() + .downcast_ref::() + .context("Storyline unknown_fields_json is not Utf8")?; + let hydrated = values + .iter() + .map(|encoded_fields| { + let Some(encoded_fields) = encoded_fields else { + return Ok(None); + }; + let mut fields: StorylineUnknownFields = serde_json::from_str(encoded_fields) + .context("decode Storyline unknown_fields_json for hydration")?; + for source in fields.sources.values_mut() { + for value in source.fields.values_mut() { + let Some(encoded) = value.as_str() else { + continue; + }; + let Some(reference) = ContentRef::parse(encoded) + .context("invalid internal unknown-field content descriptor")? + else { + continue; + }; + let object = resolved.get(&reference.content_id).with_context(|| { + format!( + "Storyline content object '{}' is missing from the committed snapshot", + reference.content_id + ) + })?; + anyhow::ensure!( + reference.logical_type == LogicalType::Json, + "unknown-field content descriptor is not JSON" + ); + anyhow::ensure!( + object.codec == reference.codec + && object.raw_length == reference.raw_length, + "Storyline content descriptor metadata mismatch for '{}'", + reference.content_id + ); + *value = serde_json::from_slice(&object.bytes).with_context(|| { + format!( + "Storyline unknown-field object '{}' is invalid JSON", + reference.content_id + ) + })?; + } + } + serde_json::to_string(&fields) + .map(Some) + .context("encode hydrated Storyline unknown_fields_json") + }) + .collect::>>()?; + let mut columns = batch.columns().to_vec(); + columns[index] = Arc::new(StringArray::from(hydrated)); + RecordBatch::try_new(batch.schema(), columns) + .context("hydrate Storyline unknown-field content batch") } pub(crate) async fn hydrate_selected_batches( @@ -779,7 +969,7 @@ fn hydrate_batch( ) })?; anyhow::ensure!( - object.raw_length == reference.raw_length, + object.codec == reference.codec && object.raw_length == reference.raw_length, "Storyline content descriptor metadata mismatch for '{}'", reference.content_id ); @@ -845,7 +1035,14 @@ async fn resolve_objects( ); anyhow::ensure!( resolved - .insert(content_id.clone(), ResolvedObject { raw_length, bytes },) + .insert( + content_id.clone(), + ResolvedObject { + codec, + raw_length, + bytes, + }, + ) .is_none(), "duplicate Storyline content object '{content_id}'" ); diff --git a/crates/persisting-pchronicle/src/store/storyline/mod.rs b/crates/persisting-pchronicle/src/store/storyline/mod.rs index 4f409e57..95d1aa09 100644 --- a/crates/persisting-pchronicle/src/store/storyline/mod.rs +++ b/crates/persisting-pchronicle/src/store/storyline/mod.rs @@ -71,15 +71,16 @@ use object_store::{Error as ObjectStoreError, ObjectStoreExt, PutMode, UpdateVer use serde::{Deserialize, Serialize}; use super::storyline_model::{ - reconstruct_storyline, split_storyline, StoryRunRow, StoryStepRow, StoryToolCallRow, - StorylineTables, STORY_RUNS_TABLE, STORY_STEPS_TABLE, STORY_TOOL_CALLS_TABLE, + reconstruct_storyline, split_storyline_with_unknown_limits, StoryRunRow, StoryStepRow, + StoryToolCallRow, StorylineTables, STORY_RUNS_TABLE, STORY_STEPS_TABLE, STORY_TOOL_CALLS_TABLE, }; +use crate::formats::unknown_fields::{compute_unknown_key_counts, validate_unknown_fields}; use crate::StorylineDocument; use self::content::{ collect_content_ids, commit_pending_content, content_columns, externalize_batches, - hydrate_batches, open_objects, prune_unreferenced_objects, PendingContent, - STORYLINE_OBJECTS_DATASET, + externalize_unknown_field_values, hydrate_batches, open_objects, prune_unreferenced_objects, + PendingContent, STORYLINE_OBJECTS_DATASET, }; use super::AtifReader; use super::{root_write_lock, LanceMaintenanceOptions, LanceMaintenanceReport}; @@ -1124,7 +1125,8 @@ impl StorylineLanceStore { _ => session_id.to_string(), }; let mut runs = HashMap::with_capacity(ids.len()); - for run in decode_run_batches(&run_batches)? { + for mut run in decode_run_batches(&run_batches)? { + run.unknown_key_counts = compute_unknown_key_counts(&run.unknown_fields)?; let key = row_key(&run.document_id, &run.session_id); if runs.insert(key.clone(), run).is_some() { anyhow::bail!("duplicate runs rows for {column} '{key}'"); @@ -1146,12 +1148,16 @@ impl StorylineLanceStore { let Some(run) = runs.remove(id) else { return Ok(None); }; - reconstruct_storyline(StorylineTables { + let story = reconstruct_storyline(StorylineTables { run, steps: steps.remove(id).unwrap_or_default(), tool_calls: tool_calls.remove(id).unwrap_or_default(), - }) - .map(Some) + })?; + validate_unknown_fields( + &story.unknown_fields, + self.content_options.unknown_field_limits(), + )?; + Ok(Some(story)) }) .collect() } @@ -1548,10 +1554,14 @@ async fn open_table_version(path: &Path, version: u64) -> Result { } fn content_column_projection(kind: StorylineTableKind) -> Vec<&'static str> { - content_columns(kind) + let mut columns = content_columns(kind) .iter() .map(|(column, _)| *column) - .collect() + .collect::>(); + if kind == StorylineTableKind::Runs { + columns.push("unknown_fields_json"); + } + columns } async fn next_storage_ordinal(paths: &StorylineTablePaths) -> Result { diff --git a/crates/persisting-pchronicle/src/store/storyline/model.rs b/crates/persisting-pchronicle/src/store/storyline/model.rs index 82079590..60df735c 100644 --- a/crates/persisting-pchronicle/src/store/storyline/model.rs +++ b/crates/persisting-pchronicle/src/store/storyline/model.rs @@ -9,10 +9,10 @@ use std::collections::{BTreeMap, HashSet}; use serde::{Deserialize, Serialize}; use serde_json::Value; -use crate::{ - FieldPresence, Result, StoryLink, StorylineDocument, StorylinePresence, StorylineToolCall, - StorylineTurn, +use crate::formats::unknown_fields::{ + validate_unknown_fields, StorylineUnknownFields, UnknownFieldLimits, UnknownKeyCounts, }; +use crate::{Result, StoryLink, StorylineDocument, StorylineToolCall, StorylineTurn}; #[cfg(feature = "lance-store")] pub const STORY_RUNS_TABLE: &str = "runs"; @@ -27,8 +27,7 @@ pub struct StoryRunRow { /// Stable per-document storage identity. Explicit ATIF `trajectory_id` /// wins; otherwise the effective `session_id` is used. pub document_id: String, - /// Stable global order inside this Storyline store. This is deliberately - /// distinct from source-container ordinals retained in `presence`. + /// Stable global order inside this Storyline store. pub storage_ordinal: i64, pub trajectory_id_explicit: bool, pub run_id: Option, @@ -46,7 +45,8 @@ pub struct StoryRunRow { pub final_metrics: Option, pub continued_trajectory_ref: Option, pub extra: Option, - pub presence: StorylinePresence, + pub unknown_fields: StorylineUnknownFields, + pub unknown_key_counts: UnknownKeyCounts, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] @@ -85,7 +85,7 @@ pub struct StoryToolCallRow { pub tool_call_id: String, pub function_name: String, pub arguments: Value, - pub result: FieldPresence, + pub result: Option, pub results: Vec, pub duration_ms: Option, pub extra: Option, @@ -118,7 +118,19 @@ fn source_call_id(result: &Value) -> Option<&str> { } pub fn split_storyline(story: &StorylineDocument) -> Result { + split_storyline_with_unknown_limits(story, UnknownFieldLimits::default()) +} + +pub(crate) fn split_storyline_with_unknown_limits( + story: &StorylineDocument, + unknown_limits: UnknownFieldLimits, +) -> Result { story.validate()?; + let counts = validate_unknown_fields(&story.unknown_fields, unknown_limits)?; + anyhow::ensure!( + counts == story.unknown_key_counts, + "storyline unknown_key_counts do not match unknown_fields" + ); let document_id = story .trajectory_id .clone() @@ -143,7 +155,8 @@ pub fn split_storyline(story: &StorylineDocument) -> Result { final_metrics: story.final_metrics.clone(), continued_trajectory_ref: story.continued_trajectory_ref.clone(), extra: story.extra.clone(), - presence: story.presence.clone(), + unknown_fields: story.unknown_fields.clone(), + unknown_key_counts: story.unknown_key_counts.clone(), }; let mut seen_calls = HashSet::new(); @@ -336,7 +349,6 @@ pub fn reconstruct_storyline(tables: StorylineTables) -> Result Result, + mut runs: Vec, steps: Vec, tool_calls: Vec, options: StorylineContentOptions, ) -> Result { let mut pending = PendingContent::default(); + for run in &mut runs { + externalize_unknown_field_values(&mut run.unknown_fields, options, &mut pending)?; + } let runs = externalize_batches( encode_rows(runs, story_runs_to_batch)?, StorylineTableKind::Runs, diff --git a/crates/persisting-pchronicle/src/store/storyline/rows.rs b/crates/persisting-pchronicle/src/store/storyline/rows.rs index 6b5193c6..9867f804 100644 --- a/crates/persisting-pchronicle/src/store/storyline/rows.rs +++ b/crates/persisting-pchronicle/src/store/storyline/rows.rs @@ -11,7 +11,6 @@ use serde::de::DeserializeOwned; use serde::Serialize; use super::super::storyline_model::{StoryRunRow, StoryStepRow, StoryToolCallRow}; -use crate::FieldPresence; fn field(name: &str, data_type: DataType, nullable: bool) -> Field { Field::new(name, data_type, nullable) @@ -38,7 +37,8 @@ pub fn story_runs_arrow_schema() -> Arc { field("final_metrics_json", DataType::Utf8, true), field("continued_trajectory_ref", DataType::Utf8, true), field("extra_json", DataType::Utf8, true), - field("presence_json", DataType::Utf8, true), + field("unknown_fields_json", DataType::Utf8, true), + field("unknown_key_counts_json", DataType::Utf8, true), ])) } @@ -83,7 +83,6 @@ pub fn story_tool_calls_arrow_schema() -> Arc { field("tool_call_id", DataType::Utf8, false), field("function_name", DataType::Utf8, false), field("arguments_json", DataType::Utf8, false), - field("result_present", DataType::Boolean, false), field("result_json", DataType::Utf8, true), field("results_json", DataType::Utf8, false), field("duration_ms", DataType::Int64, true), @@ -188,9 +187,22 @@ pub fn story_runs_to_batch(rows: &[StoryRunRow]) -> Result { .map(|r| opt_json(&r.extra)) .collect::>>()?, )), - Arc::new(req_utf8_owned( + Arc::new(opt_utf8_owned( + rows.iter() + .map(|r| { + (!r.unknown_fields.is_empty()) + .then(|| json(&r.unknown_fields)) + .transpose() + }) + .collect::>>()?, + )), + Arc::new(opt_utf8_owned( rows.iter() - .map(|r| json(&r.presence)) + .map(|r| { + (!r.unknown_key_counts.is_empty()) + .then(|| json(&r.unknown_key_counts)) + .transpose() + }) .collect::>>()?, )), ], @@ -279,17 +291,9 @@ pub fn story_tool_calls_to_batch(rows: &[StoryToolCallRow]) -> Result>()?, )), - Arc::new(BooleanArray::from( - rows.iter() - .map(|r| !r.result.is_missing()) - .collect::>(), - )), Arc::new(opt_utf8_owned( rows.iter() - .map(|r| match &r.result { - FieldPresence::Value(value) => json(value).map(Some), - FieldPresence::Missing | FieldPresence::Null => Ok(None), - }) + .map(|r| r.result.as_ref().map(json).transpose()) .collect::>>()?, )), Arc::new(req_utf8_owned( @@ -433,8 +437,14 @@ pub fn story_runs_from_batch(batch: &RecordBatch) -> Result> { final_metrics: optional_json_at(batch, "final_metrics_json", row)?, continued_trajectory_ref: string_at(batch, "continued_trajectory_ref", row)?, extra: optional_json_at(batch, "extra_json", row)?, - presence: optional_json_if_present(batch, "presence_json", row)? + unknown_fields: optional_json_if_present(batch, "unknown_fields_json", row)? .unwrap_or_default(), + unknown_key_counts: optional_json_if_present( + batch, + "unknown_key_counts_json", + row, + )? + .unwrap_or_default(), }) }) .collect() @@ -489,14 +499,9 @@ pub fn story_tool_calls_from_batch(batch: &RecordBatch) -> Result FieldPresence::Value(parse_json(value, "result_json")?), - None => FieldPresence::Null, - } - } else { - FieldPresence::Missing - }, + result: string_at(batch, "result_json", row)? + .map(|value| parse_json(value, "result_json")) + .transpose()?, results: parse_json( required_string_at(batch, "results_json", row)?, "results_json", @@ -514,9 +519,9 @@ mod tests { #[test] fn empty_batches_keep_all_three_schemas() { - assert_eq!(story_runs_to_batch(&[]).unwrap().num_columns(), 20); + assert_eq!(story_runs_to_batch(&[]).unwrap().num_columns(), 21); assert_eq!(story_steps_to_batch(&[]).unwrap().num_columns(), 20); - assert_eq!(story_tool_calls_to_batch(&[]).unwrap().num_columns(), 13); + assert_eq!(story_tool_calls_to_batch(&[]).unwrap().num_columns(), 12); } #[test] @@ -598,7 +603,7 @@ mod tests { tool_call_id: "c".into(), function_name: "lookup".into(), arguments: serde_json::json!({"q": "x"}), - result: FieldPresence::Null, + result: None, results: vec![serde_json::json!({"source_call_id": "c", "content": "y"})], duration_ms: Some(8), extra: None, diff --git a/crates/persisting-pchronicle/src/store/storyline/tests.rs b/crates/persisting-pchronicle/src/store/storyline/tests.rs index 63a52e5e..75bf776e 100644 --- a/crates/persisting-pchronicle/src/store/storyline/tests.rs +++ b/crates/persisting-pchronicle/src/store/storyline/tests.rs @@ -110,7 +110,8 @@ fn story(session_id: &str) -> StorylineDocument { final_metrics: None, continued_trajectory_ref: None, extra: None, - presence: Default::default(), + unknown_fields: Default::default(), + unknown_key_counts: Default::default(), turns: vec![ StorylineTurn { id: 1, @@ -161,6 +162,198 @@ fn story(session_id: &str) -> StorylineDocument { } } +#[test] +fn unknown_field_limit_options_reject_zero_and_unbounded() { + for options in [ + StorylineContentOptions { + max_unknown_fields: 0, + ..Default::default() + }, + StorylineContentOptions { + max_unknown_bytes: 0, + ..Default::default() + }, + StorylineContentOptions { + max_unknown_fields: usize::MAX, + ..Default::default() + }, + StorylineContentOptions { + max_unknown_bytes: usize::MAX, + ..Default::default() + }, + ] { + assert!(options.validate().is_err()); + } +} + +#[tokio::test] +async fn configured_unknown_field_limit_reports_actual_and_limit() { + let temporary = tempfile::tempdir().unwrap(); + let store = StorylineLanceStore::open_with_content_options( + temporary.path(), + StorylineContentOptions { + max_unknown_fields: 1, + ..Default::default() + }, + ) + .await + .unwrap(); + let mut oversized = story("unknown-count-limit"); + oversized + .unknown_fields + .insert("atif", "source", "/one", serde_json::json!(1)) + .unwrap(); + oversized + .unknown_fields + .insert("atif", "source", "/two", serde_json::json!(2)) + .unwrap(); + oversized.refresh_unknown_key_counts().unwrap(); + + let error = store.replace_storyline(&oversized).await.unwrap_err(); + assert!( + error + .to_string() + .contains("unknown field count 2 exceeds configured limit 1"), + "{error:#}" + ); +} + +#[tokio::test] +async fn repeated_unknown_value_is_stored_once() { + let large = serde_json::json!({ + "payload": "x".repeat(DEFAULT_CONTENT_OFFLOAD_THRESHOLD) + }); + let mut first = story("unknown-first"); + first + .unknown_fields + .insert("actf", "task-1", "/shared", large.clone()) + .unwrap(); + first.refresh_unknown_key_counts().unwrap(); + let mut second = story("unknown-second"); + second + .unknown_fields + .insert("actf", "task-1", "/shared", large.clone()) + .unwrap(); + second.refresh_unknown_key_counts().unwrap(); + let temporary = tempfile::tempdir().unwrap(); + let store = StorylineLanceStore::open(temporary.path()).await.unwrap(); + store.replace_storylines(&[first, second]).await.unwrap(); + + let paths = store.current_table_paths().await.unwrap().unwrap(); + let objects = open_objects(&paths.objects, paths.objects_version) + .await + .unwrap(); + assert_eq!(objects.count_rows(None).await.unwrap(), 1); + let hydrated = store + .get_storyline_full("unknown-first") + .await + .unwrap() + .unwrap(); + assert_eq!( + hydrated.unknown_fields.sources["actf"].fields["/shared"], + large + ); +} + +#[tokio::test] +async fn unknown_content_ref_magic_string_round_trips_as_literal() { + let literal = format!("{CONTENT_REF_MAGIC}user-controlled-not-a-descriptor"); + let mut expected = story("unknown-magic"); + expected + .unknown_fields + .insert("atif", "source", "/literal", serde_json::json!(literal)) + .unwrap(); + expected.refresh_unknown_key_counts().unwrap(); + let temporary = tempfile::tempdir().unwrap(); + let store = StorylineLanceStore::open_with_content_options( + temporary.path(), + StorylineContentOptions { + offload_threshold: usize::MAX, + ..Default::default() + }, + ) + .await + .unwrap(); + store.replace_storyline(&expected).await.unwrap(); + assert_eq!( + store.get_storyline_full("unknown-magic").await.unwrap(), + Some(expected) + ); +} + +#[tokio::test] +async fn logical_unknown_limit_rejects_compressible_value_before_offload() { + let mut oversized = story("logical-unknown-limit"); + oversized + .unknown_fields + .insert( + "atif", + "source", + "/payload", + serde_json::json!("x".repeat(crate::model::DEFAULT_MAX_UNKNOWN_BYTES + 1)), + ) + .unwrap(); + oversized.refresh_unknown_key_counts().unwrap(); + let temporary = tempfile::tempdir().unwrap(); + let store = StorylineLanceStore::open(temporary.path()).await.unwrap(); + + let error = store.replace_storyline(&oversized).await.unwrap_err(); + assert!( + error.to_string().contains("unknown field byte size") + && error + .to_string() + .contains("exceeds configured limit 1048576"), + "{error:#}" + ); + assert!(store.current_table_paths().await.unwrap().is_none()); +} + +#[tokio::test] +async fn logical_unknown_limit_rejects_hydrated_value_on_read() { + let temporary = tempfile::tempdir().unwrap(); + let writer = StorylineLanceStore::open_with_content_options( + temporary.path(), + StorylineContentOptions { + offload_threshold: 1, + max_unknown_bytes: 4096, + ..Default::default() + }, + ) + .await + .unwrap(); + let mut stored = story("logical-read-limit"); + stored + .unknown_fields + .insert( + "atif", + "source", + "/payload", + serde_json::json!("x".repeat(128)), + ) + .unwrap(); + stored.refresh_unknown_key_counts().unwrap(); + writer.replace_storyline(&stored).await.unwrap(); + + let reader = StorylineLanceStore::open_with_content_options( + temporary.path(), + StorylineContentOptions { + max_unknown_bytes: 16, + ..Default::default() + }, + ) + .await + .unwrap(); + let error = reader + .get_storyline_full("logical-read-limit") + .await + .unwrap_err(); + assert!( + error.to_string().contains("unknown field byte size") + && error.to_string().contains("exceeds configured limit 16"), + "{error:#}" + ); +} + #[tokio::test] async fn persists_three_tables_and_round_trips_storyline() { let dir = tempfile::tempdir().unwrap(); @@ -1017,6 +1210,9 @@ async fn atif_stream_preserves_nested_documents_with_shared_session() { let dir = tempfile::tempdir().unwrap(); let store = StorylineLanceStore::open(dir.path()).await.unwrap(); + let canonical_stories = crate::convert::atif_collection_to_storylines(input.clone()).unwrap(); + let canonical = crate::convert::storylines_to_atif(&canonical_stories).unwrap(); + let report = store.import_atif_stream(file.path()).await.unwrap(); assert_eq!(report.storylines, 2); let stories = store @@ -1027,7 +1223,7 @@ async fn atif_stream_preserves_nested_documents_with_shared_session() { .collect::>>() .unwrap(); let rebuilt = crate::convert::storylines_to_atif(&stories).unwrap(); - assert_eq!(serde_json::to_value(&rebuilt[0]).unwrap(), input); + assert_eq!(rebuilt, canonical); } #[tokio::test] diff --git a/crates/persisting-pchronicle/src/tests.rs b/crates/persisting-pchronicle/src/tests.rs index 515e78b6..fc93a7e0 100644 --- a/crates/persisting-pchronicle/src/tests.rs +++ b/crates/persisting-pchronicle/src/tests.rs @@ -46,7 +46,9 @@ fn from_storyline(format: TestFormat, story: &crate::StorylineDocument) -> crate TestFormat::CanonicalEvent => Err(lance_only_error()), TestFormat::AgenticMd => crate::document::encode_agenticmd(story), TestFormat::OpenaiMsg => Ok(serde_json::to_string_pretty( - &crate::formats::synthesize_openai_msg_corpus(std::slice::from_ref(story))?, + &crate::formats::openai_corpus::synthesize_openai_msg_corpus(std::slice::from_ref( + story, + ))?, )?), TestFormat::Atif => Ok(serde_json::to_string_pretty( &crate::convert::storyline_to_atif(story)?, @@ -72,75 +74,81 @@ fn lance_only_error() -> anyhow::Error { fn sample_traj() -> AtifTrajectory { AtifTrajectory { schema_version: "ATIF-v1.7".into(), - session_id: crate::FieldPresence::Value("sess-1".into()), - trajectory_id: crate::FieldPresence::Value("traj-1".into()), + unknown: Default::default(), + session_id: Some("sess-1".into()), + trajectory_id: Some("traj-1".into()), agent: AtifAgent { name: "harbor-agent".into(), version: "1.0.0".into(), - model_name: crate::FieldPresence::Value("gemini-2.5-flash".into()), - tool_definitions: crate::FieldPresence::Missing, - extra: crate::FieldPresence::Missing, + unknown: Default::default(), + model_name: Some("gemini-2.5-flash".into()), + tool_definitions: None, + extra: None, }, - notes: crate::FieldPresence::Value("unit test".into()), - final_metrics: crate::FieldPresence::Value(json!({"total_steps": 2})), - continued_trajectory_ref: crate::FieldPresence::Missing, - extra: crate::FieldPresence::Missing, - subagent_trajectories: crate::FieldPresence::Missing, + notes: Some("unit test".into()), + final_metrics: Some(json!({"total_steps": 2})), + continued_trajectory_ref: None, + extra: None, + subagent_trajectories: None, steps: vec![ AtifStep { step_id: 1, - timestamp: crate::FieldPresence::Value("2025-10-11T10:30:00Z".into()), + unknown: Default::default(), + timestamp: Some("2025-10-11T10:30:00Z".into()), source: "user".into(), - model_name: crate::FieldPresence::Missing, - reasoning_effort: crate::FieldPresence::Missing, + model_name: None, + reasoning_effort: None, message: json!("What is the price of GOOGL?"), - reasoning_content: crate::FieldPresence::Missing, - tool_calls: crate::FieldPresence::Missing, - observation: crate::FieldPresence::Missing, - metrics: crate::FieldPresence::Missing, - extra: crate::FieldPresence::Missing, - llm_call_count: crate::FieldPresence::Missing, - is_copied_context: crate::FieldPresence::Missing, + reasoning_content: None, + tool_calls: None, + observation: None, + metrics: None, + extra: None, + llm_call_count: None, + is_copied_context: None, }, AtifStep { step_id: 2, - timestamp: crate::FieldPresence::Value("2025-10-11T10:30:02Z".into()), + unknown: Default::default(), + timestamp: Some("2025-10-11T10:30:02Z".into()), source: "agent".into(), - model_name: crate::FieldPresence::Value("gemini-2.5-flash".into()), - reasoning_effort: crate::FieldPresence::Value(json!("medium")), + model_name: Some("gemini-2.5-flash".into()), + reasoning_effort: Some(json!("medium")), message: json!("I will search."), - reasoning_content: crate::FieldPresence::Value("Need price and volume.".into()), - tool_calls: crate::FieldPresence::Value(vec![ + reasoning_content: Some("Need price and volume.".into()), + tool_calls: Some(vec![ AtifToolCall { tool_call_id: "call_price_1".into(), + unknown: Default::default(), function_name: "financial_search".into(), arguments: json!({"ticker":"GOOGL","metric":"price"}), - result: crate::FieldPresence::Value(json!({"price": 185.35})), - extra: crate::FieldPresence::Value(json!({"duration_ms": 42})), + result: Some(json!({"price": 185.35})), + extra: Some(json!({"duration_ms": 42})), }, AtifToolCall { tool_call_id: "call_volume_2".into(), + unknown: Default::default(), function_name: "financial_search".into(), arguments: json!({"ticker":"GOOGL","metric":"volume"}), - result: crate::FieldPresence::Missing, - extra: crate::FieldPresence::Value(json!({"duration_ms": 37})), + result: None, + extra: Some(json!({"duration_ms": 37})), }, ]), - observation: crate::FieldPresence::Value(AtifObservation { + observation: Some(AtifObservation { results: vec![ json!({"source_call_id":"call_price_1","content":"$185.35"}), json!({"source_call_id":"call_volume_2","content":"1.5M"}), ], }), - metrics: crate::FieldPresence::Value(json!({ + metrics: Some(json!({ "prompt_tokens": 520, "completion_tokens": 80, "latency_ms": 1850, "ttft_ms": 210 })), - extra: crate::FieldPresence::Missing, - llm_call_count: crate::FieldPresence::Value(1), - is_copied_context: crate::FieldPresence::Missing, + extra: None, + llm_call_count: Some(1), + is_copied_context: None, }, ], } @@ -173,7 +181,7 @@ fn atif_storyline_hub_roundtrip() { assert_eq!(back.steps[1].tool_calls.as_ref().unwrap().len(), 2); assert_eq!( back.steps[1].tool_calls.as_ref().unwrap()[0].result, - crate::FieldPresence::Value(serde_json::json!({"price": 185.35})) + Some(serde_json::json!({"price": 185.35})) ); assert_eq!( back.steps[1] @@ -634,7 +642,7 @@ fn convert_openai_msg_storyline_roundtrip_messages() { let back = convert(TestFormat::Storyline, TestFormat::OpenaiMsg, &story).unwrap(); let doc: serde_json::Value = serde_json::from_str(&back).unwrap(); - let rows = doc.as_array().unwrap(); + let rows = doc["session_steps"].as_array().unwrap(); assert_eq!(rows[0]["session_id"], "s-om"); assert!(!rows.is_empty()); } @@ -746,7 +754,8 @@ fn storyline_to_events_assigns_call_id_for_paired_turns() { final_metrics: None, continued_trajectory_ref: None, extra: None, - presence: Default::default(), + unknown_fields: Default::default(), + unknown_key_counts: Default::default(), turns: vec![ StorylineTurn { id: 1, diff --git a/crates/persisting-pchronicle/tests/README.md b/crates/persisting-pchronicle/tests/README.md index c01c52b3..b808efc8 100644 --- a/crates/persisting-pchronicle/tests/README.md +++ b/crates/persisting-pchronicle/tests/README.md @@ -19,6 +19,7 @@ crate 级门面行为（格式往返、detect、frontmatter 解析）集中在 ` |---|---| | `atif_lance_corpus.rs` | 独立 ATIF corpus（`tests/fixtures/atif/`，8 条 10–20 step 确定性轨迹）的格式转换、Storyline 三表落盘与空间占用 | | `capture_fixture_corpus.rs` | 复用 `persisting-gateway/tests/fixtures` 的跨组件语料：AgenticMD golden 往返、request/response/provider snapshot/SSE 文本的无损往返（设置最小样本数，防止 fixture 缩减后静默通过） | +| `conversion_semantics.rs` | ATIF、ACTF、OpenAI 的三条直接 Lance 往返与六条有向跨格式转换；对象字段的 `null`/missing 按语义等价比较，数组中的 `null` 保持显著，unknown field 的 JSON Pointer、值与出现次数精确校验 | | `direct_file_query.rs` | 直接对文件/目录的只读查询（不经过导入） | | `import_roundtrip_fixtures.rs` | `tests/fixtures/import_roundtrip/` 中 OpenAI corpus / ACTF 经三表 Lance 的无损恢复 | | `langfuse_backend_faults.rs` | 存储后端故障语义（追加失败、重复、未知错误分类） | @@ -26,6 +27,7 @@ crate 级门面行为（格式往返、detect、frontmatter 解析）集中在 ` | `query_engine.rs` | `ChronicleQueryEngine` 统一 SQL 面（Lance/ATIF/OpenAI/ACTF、`_file_`、只读门禁） | | `s3_storage.rs` | 真实 S3/MinIO 对象存储契约（默认忽略，见下） | | `search_integration.rs` | Search 索引与检索（Cargo.toml 显式声明的 `[[test]]`） | +| `unknown_fields_roundtrip.rs` | unknown-fields envelope 在 ATIF、ACTF、OpenAI 跨格式链路中的 JSON Pointer、值与计数保真 | 真实 S3/MinIO 契约测试在隔离的测试前缀下显式运行： diff --git a/crates/persisting-pchronicle/tests/atif_lance_corpus.rs b/crates/persisting-pchronicle/tests/atif_lance_corpus.rs index 4e083352..2dff8e74 100644 --- a/crates/persisting-pchronicle/tests/atif_lance_corpus.rs +++ b/crates/persisting-pchronicle/tests/atif_lance_corpus.rs @@ -168,10 +168,10 @@ async fn corpus_round_trips_through_three_lance_tables() -> Result<()> { } #[tokio::test] -async fn atif_null_and_tool_result_presence_round_trip_through_lance() -> Result<()> { +async fn atif_null_and_missing_canonicalization_is_stable_through_lance() -> Result<()> { let input = serde_json::json!({ "schema_version": "ATIF-v1.7", - "trajectory_id": "presence-trajectory", + "trajectory_id": "canonical-trajectory", "agent": { "name": "agent-1", "version": "1", @@ -216,6 +216,8 @@ async fn atif_null_and_tool_result_presence_round_trip_through_lance() -> Result "subagent_trajectories": null }); let story = into_storyline(TestFormat::Atif, &input.to_string())?; + let expected: serde_json::Value = + serde_json::from_str(&from_storyline(TestFormat::Atif, &story)?)?; let dir = tempfile::tempdir()?; let store = StorylineLanceStore::open(dir.path()).await?; store.replace_storyline(&story).await?; @@ -223,10 +225,10 @@ async fn atif_null_and_tool_result_presence_round_trip_through_lance() -> Result let restored = store .get_storyline_full(&story.session_id) .await? - .context("missing presence Storyline after Lance write")?; + .context("missing canonical Storyline after Lance write")?; let output: serde_json::Value = serde_json::from_str(&from_storyline(TestFormat::Atif, &restored)?)?; - assert_eq!(output, input); + assert_eq!(output, expected); Ok(()) } diff --git a/crates/persisting-pchronicle/tests/conversion_semantics.rs b/crates/persisting-pchronicle/tests/conversion_semantics.rs new file mode 100644 index 00000000..0a4cc9df --- /dev/null +++ b/crates/persisting-pchronicle/tests/conversion_semantics.rs @@ -0,0 +1,421 @@ +use std::path::Path; + +use anyhow::{Context, Result}; +use persisting_pchronicle::document::{ + decode_json_storylines, encode_json_storylines, DocumentFormat, +}; +use persisting_pchronicle::model::StorylineDocument; +use persisting_pchronicle::storage::StorylineLanceStore; +use serde_json::{json, Map, Value}; + +const JSON_FORMATS: [DocumentFormat; 3] = [ + DocumentFormat::Atif, + DocumentFormat::Actf, + DocumentFormat::OpenaiMsg, +]; + +struct FormatCase { + name: &'static str, + format: DocumentFormat, + relative_path: &'static str, + input: Value, +} + +fn canonical_semantic_json(value: Value) -> Value { + match value { + Value::Object(fields) => Value::Object( + fields + .into_iter() + .filter(|(_, value)| !value.is_null()) + .map(|(key, value)| (key, canonical_semantic_json(value))) + .collect::>(), + ), + Value::Array(values) => { + Value::Array(values.into_iter().map(canonical_semantic_json).collect()) + } + value => value, + } +} + +fn assert_semantic_json_eq(actual: &Value, expected: &Value) { + assert_eq!( + canonical_semantic_json(actual.clone()), + canonical_semantic_json(expected.clone()) + ); +} + +fn format_cases() -> Vec { + vec![ + FormatCase { + name: "atif", + format: DocumentFormat::Atif, + relative_path: "semantic.atif.json", + input: json!({ + "schema_version": "ATIF-v1.7", + "trajectory_id": "semantic-atif", + "session_id": null, + "agent": { + "name": "agent", + "version": "1", + "model_name": null, + "vendor_agent": {"enabled": true} + }, + "steps": [ + { + "step_id": 1, + "source": "user", + "message": "inspect", + "reasoning_content": null, + "vendor_step": {"ordinal": 1}, + "vendor_null": null + }, + { + "step_id": 2, + "source": "agent", + "message": "done", + "tool_calls": [{ + "tool_call_id": "call-1", + "function_name": "inspect", + "arguments": {"path": "/tmp"}, + "result": null, + "vendor_call": [1, 2, 3] + }], + "observation": { + "results": [{"source_call_id": "call-1", "content": "ok"}] + }, + "vendor_step": {"ordinal": 2} + } + ], + "vendor/root": {"tilde~key": 7} + }), + }, + FormatCase { + name: "actf", + format: DocumentFormat::Actf, + relative_path: "semantic.actf.json", + input: json!({ + "task_id": "semantic-actf", + "category": "regression", + "k": 1, + "correct": true, + "attempts_tried": 1, + "solved_at": null, + "attempts": { + "1": { + "correct": true, + "final_answer": null, + "ground_truth": "done", + "trajectory": { + "schema_version": "ACTF_v1.0", + "steps": [ + { + "step_id": 1, + "assistant_content": { + "content": "working", + "reasoning_content": "inspect", + "tool_calls": [] + }, + "metric": { + "prompt_tokens_len": null, + "completion_tokens_len": null, + "llm_infer_ms": null, + "env_action_ms": null, + "stop_reason": null + }, + "system_prompt": "system", + "user_content": "inspect", + "tools": [], + "observation": [], + "started_at": "2026-08-20T00:00:00Z", + "finished_at": "2026-08-20T00:00:01Z", + "vendor_step": {"ordinal": 1}, + "vendor_null": null + }, + { + "step_id": 2, + "assistant_content": { + "content": "done", + "reasoning_content": "complete", + "tool_calls": [{ + "id": "call-1", + "type": "tool_use", + "name": "inspect", + "input": {"path": "/tmp"}, + "vendor_call": true + }] + }, + "metric": { + "prompt_tokens_len": 10, + "completion_tokens_len": 4, + "llm_infer_ms": 12, + "env_action_ms": 3, + "stop_reason": "stop" + }, + "system_prompt": "system", + "user_content": "", + "tools": [{ + "id": "call-1", + "type": "tool_use", + "name": "inspect", + "input": {"path": "/tmp"}, + "vendor_call": true + }], + "observation": [{ + "id": "call-1", + "type": "tool_result", + "content": "ok" + }], + "started_at": "2026-08-20T00:00:01Z", + "finished_at": "2026-08-20T00:00:02Z", + "vendor_step": {"ordinal": 2} + } + ], + "started_at": "2026-08-20T00:00:00Z", + "finished_at": "2026-08-20T00:00:02Z", + "vendor_trajectory": "kept" + }, + "status": "completed", + "score": null, + "error": "", + "artifacts": {}, + "extra": {}, + "analysis_result": {}, + "meta": {"suite": "semantic"}, + "vendor_attempt": ["kept"] + } + }, + "vendor/root": {"tilde~key": 7} + }), + }, + FormatCase { + name: "openai", + format: DocumentFormat::OpenaiMsg, + relative_path: "semantic.openai.json", + input: json!({ + "vendor/root": {"tilde~key": 7}, + "root_null": null, + "session_steps": [{ + "session_id": "semantic-openai", + "step_id": 1, + "created_at": 1787184000, + "messages": [{ + "role": "user", + "content": "inspect", + "vendor_message": {"ordinal": 1} + }], + "response": { + "role": "assistant", + "content": "done", + "name": null, + "tool_calls": null, + "vendor_response": [1, 2] + }, + "agent_model": "model", + "reward": null, + "vendor_row": {"enabled": true}, + "vendor_null": null + }] + }), + }, + ] +} + +async fn persist_and_restore(stories: &[StorylineDocument]) -> Result> { + let temporary = tempfile::tempdir()?; + let store = StorylineLanceStore::open(temporary.path()).await?; + store.replace_storylines(stories).await?; + let document_ids = stories + .iter() + .map(StorylineDocument::document_id) + .map(str::to_owned) + .collect::>(); + store + .get_storylines_by_document_ids(&document_ids) + .await? + .into_iter() + .map(|story| story.context("Lance roundtrip lost a Storyline")) + .collect() +} + +fn assert_unknown_contract(format: DocumentFormat, stories: &[StorylineDocument]) { + assert_eq!(stories.len(), 1); + let story = &stories[0]; + match format { + DocumentFormat::Atif => { + let source = &story.unknown_fields.sources["atif"]; + assert_eq!(source.fields["/vendor~1root"], json!({"tilde~key": 7})); + assert_eq!(source.fields["/steps/0/vendor_null"], Value::Null); + assert_eq!(story.unknown_key_counts["atif"]["/steps/*/vendor_step"], 2); + assert_eq!( + source.fields["/steps/1/tool_calls/0/vendor_call"], + json!([1, 2, 3]) + ); + } + DocumentFormat::Actf => { + let source = &story.unknown_fields.sources["actf"]; + assert_eq!(source.fields["/vendor~1root"], json!({"tilde~key": 7})); + assert_eq!( + source.fields["/attempts/1/trajectory/steps/0/vendor_null"], + Value::Null + ); + assert_eq!( + story.unknown_key_counts["actf"]["/attempts/1/trajectory/steps/*/vendor_step"], + 2 + ); + assert_eq!(source.fields["/attempts/1/vendor_attempt"], json!(["kept"])); + } + DocumentFormat::OpenaiMsg => { + let source = &story.unknown_fields.sources["openai-msg"]; + assert_eq!(source.fields["/vendor~1root"], json!({"tilde~key": 7})); + assert_eq!(source.fields["/session_steps/0/vendor_null"], Value::Null); + assert_eq!( + story.unknown_key_counts["openai-msg"]["/session_steps/*/vendor_row"], + 1 + ); + assert_eq!( + source.fields["/session_steps/0/messages/0/vendor_message"], + json!({"ordinal": 1}) + ); + } + unsupported => panic!("unexpected JSON format {unsupported}"), + } +} + +fn common_storyline_semantics(stories: &[StorylineDocument]) -> Value { + Value::Array( + stories + .iter() + .map(|story| { + json!({ + "session_id": story.session_id, + "turns": story.turns.iter().map(|turn| { + let tool_calls = turn.tool_calls.as_deref().unwrap_or_default().iter() + .map(|call| json!({ + "id": call.tool_call_id, + "name": call.function_name, + "arguments": call.arguments, + })) + .collect::>(); + let observation = turn.observation.as_ref() + .and_then(|value| value.get("results")) + .and_then(Value::as_array) + .map(|results| results.iter().map(|result| json!({ + "source_call_id": result.get("source_call_id"), + "content": result.get("content"), + })).collect::>()) + .unwrap_or_default(); + json!({ + "id": turn.id, + "message": turn.message, + "tool_calls": tool_calls, + "observation": observation, + }) + }).collect::>() + }) + }) + .collect(), + ) +} + +fn assert_common_storyline_semantics(actual: &[StorylineDocument], expected: &[StorylineDocument]) { + assert_semantic_json_eq( + &common_storyline_semantics(actual), + &common_storyline_semantics(expected), + ); +} + +#[test] +fn semantic_json_treats_null_object_members_as_missing() { + assert_semantic_json_eq( + &json!({"root": null, "nested": {"value": null}}), + &json!({"nested": {}}), + ); +} + +#[test] +fn semantic_json_keeps_array_nulls_and_non_null_values_significant() { + assert_ne!( + canonical_semantic_json(json!([null])), + canonical_semantic_json(json!([])) + ); + assert_ne!( + canonical_semantic_json(json!({"value": 1})), + canonical_semantic_json(json!({"value": 2})) + ); +} + +#[tokio::test] +async fn direct_formats_survive_lance_semantically() -> Result<()> { + for case in format_cases() { + let stories = decode_json_storylines( + case.format, + &case.input.to_string(), + Path::new(case.relative_path), + ) + .with_context(|| format!("decode {} source", case.name))?; + assert_unknown_contract(case.format, &stories); + + let expected = encode_json_storylines(case.format, &stories) + .with_context(|| format!("canonicalize {} source", case.name))?; + if case.format == DocumentFormat::Atif { + assert_eq!(expected["session_id"], "semantic-atif"); + assert!(expected["steps"][1]["tool_calls"][0] + .as_object() + .is_some_and(|call| !call.contains_key("result"))); + } + + let restored = persist_and_restore(&stories) + .await + .with_context(|| format!("persist {} Storylines", case.name))?; + assert_unknown_contract(case.format, &restored); + let actual = encode_json_storylines(case.format, &restored) + .with_context(|| format!("restore {} source", case.name))?; + assert_semantic_json_eq(&actual, &expected); + } + Ok(()) +} + +fn bridge_path(format: DocumentFormat) -> &'static str { + match format { + DocumentFormat::Atif => "bridge.atif.json", + DocumentFormat::Actf => "bridge.actf.json", + DocumentFormat::OpenaiMsg => "bridge.openai.json", + unsupported => panic!("unexpected JSON format {unsupported}"), + } +} + +#[tokio::test] +async fn every_directed_cross_format_hop_is_semantically_stable_through_lance() -> Result<()> { + for source in format_cases() { + let source_stories = decode_json_storylines( + source.format, + &source.input.to_string(), + source.relative_path, + ) + .with_context(|| format!("decode {} source", source.name))?; + for bridge in JSON_FORMATS + .into_iter() + .filter(|format| *format != source.format) + { + let edge = format!("{} -> {bridge}", source.name); + let bridged = encode_json_storylines(bridge, &source_stories) + .with_context(|| format!("encode {edge}"))?; + let bridge_stories = + decode_json_storylines(bridge, &bridged.to_string(), bridge_path(bridge)) + .with_context(|| format!("decode {edge}"))?; + assert_unknown_contract(source.format, &bridge_stories); + assert_common_storyline_semantics(&bridge_stories, &source_stories); + + let restored = persist_and_restore(&bridge_stories) + .await + .with_context(|| format!("persist {edge}"))?; + assert_eq!(restored, bridge_stories, "Storyline changed across {edge}"); + assert_unknown_contract(source.format, &restored); + assert_common_storyline_semantics(&restored, &source_stories); + let actual = encode_json_storylines(bridge, &restored) + .with_context(|| format!("restore {edge} bridge"))?; + assert_semantic_json_eq(&actual, &bridged); + } + } + Ok(()) +} diff --git a/crates/persisting-pchronicle/tests/document_source.rs b/crates/persisting-pchronicle/tests/document_source.rs index 09cb00df..19f0501d 100644 --- a/crates/persisting-pchronicle/tests/document_source.rs +++ b/crates/persisting-pchronicle/tests/document_source.rs @@ -141,11 +141,19 @@ async fn opens_all_six_formats_and_reports_true_capabilities() -> Result<()> { &fixture("tests/fixtures/import_roundtrip/cybergym_0729001_trimmed.json"), ) .await?; - assert_storyline_tables( + let actf = open_document( DocumentFormat::Actf, &fixture("tests/fixtures/import_roundtrip/make-doom-for-mips_trimmed.actf.json"), ) .await?; + assert_eq!( + actf.register_datafusion(&SessionContext::new())?, + QueryTables::Storyline + ); + let actf_caps = actf.capabilities(); + assert_eq!(actf_caps.filter_pushdown, FilterPushdown::Inexact); + assert!(actf_caps.streaming_decode); + assert!(!actf.project_storylines().await?.is_empty()); Ok(()) } @@ -178,27 +186,34 @@ async fn materialization_budget_fails_closed_but_callback_visits_the_complete_st } #[tokio::test] -async fn file_document_callback_enforces_the_provider_file_budget() -> Result<()> { +async fn file_document_callbacks_enforce_the_provider_file_budget() -> Result<()> { let input = tempfile::NamedTempFile::with_suffix(".json")?; input.as_file().set_len(1024 * 1024 * 1024)?; - let source = open_document(DocumentFormat::OpenaiMsg, input.path()).await?; - - let error = source.for_each_storyline(|_| Ok(())).await.unwrap_err(); - - assert!(error.to_string().contains("exceeding max_file_bytes")); + for format in [ + DocumentFormat::Atif, + DocumentFormat::Actf, + DocumentFormat::OpenaiMsg, + ] { + let source = open_document(format, input.path()).await?; + let error = source.for_each_storyline(|_| Ok(())).await.unwrap_err(); + assert!( + error.to_string().contains("exceeding max_file_bytes"), + "{format}: {error:#}" + ); + } Ok(()) } #[tokio::test] -async fn atif_document_source_preserves_singleton_array_shape() -> Result<()> { +async fn atif_document_source_canonicalizes_singleton_array() -> Result<()> { let input = tempfile::NamedTempFile::with_suffix(".json")?; - let expected = json!([{ + let trajectory = json!({ "schema_version": "ATIF-v1.7", "trajectory_id": "one", "agent": {"name": "agent", "version": "1"}, "steps": [] - }]); - std::fs::write(input.path(), expected.to_string())?; + }); + std::fs::write(input.path(), json!([trajectory]).to_string())?; let stories = open_document(DocumentFormat::Atif, input.path()) .await? @@ -206,7 +221,13 @@ async fn atif_document_source_preserves_singleton_array_shape() -> Result<()> { .await?; assert_eq!( encode_json_storylines(DocumentFormat::Atif, &stories)?, - expected + json!({ + "schema_version": "ATIF-v1.7", + "trajectory_id": "one", + "session_id": "one", + "agent": {"name": "agent", "version": "1"}, + "steps": [] + }) ); Ok(()) } diff --git a/crates/persisting-pchronicle/tests/import_roundtrip_fixtures.rs b/crates/persisting-pchronicle/tests/import_roundtrip_fixtures.rs index b8fb463b..f8cf16c1 100644 --- a/crates/persisting-pchronicle/tests/import_roundtrip_fixtures.rs +++ b/crates/persisting-pchronicle/tests/import_roundtrip_fixtures.rs @@ -14,14 +14,12 @@ fn fixture(name: &str) -> PathBuf { async fn assert_openai_fixture_roundtrip(name: &str, expected_sessions: usize) -> Result<()> { let path = fixture(name); - let expected: serde_json::Value = serde_json::from_slice( - &std::fs::read(&path).with_context(|| format!("read fixture {}", path.display()))?, - )?; let stories = open_document(DocumentFormat::OpenaiMsg, &path) .await? .project_storylines() .await?; assert_eq!(stories.len(), expected_sessions); + let expected = encode_json_storylines(DocumentFormat::OpenaiMsg, &stories)?; let temporary = tempfile::tempdir()?; let store = StorylineLanceStore::open(temporary.path()).await?; @@ -48,8 +46,8 @@ async fn assert_actf_fixture_roundtrip(name: &str) -> Result<()> { let path = fixture(name); let raw = std::fs::read_to_string(&path) .with_context(|| format!("read fixture {}", path.display()))?; - let expected: serde_json::Value = serde_json::from_str(&raw)?; let stories = decode_json_storylines(DocumentFormat::Actf, &raw, name)?; + let expected = encode_json_storylines(DocumentFormat::Actf, &stories)?; let temporary = tempfile::tempdir()?; let store = StorylineLanceStore::open(temporary.path()).await?; diff --git a/crates/persisting-pchronicle/tests/query_engine.rs b/crates/persisting-pchronicle/tests/query_engine.rs index 07d3760f..56304a01 100644 --- a/crates/persisting-pchronicle/tests/query_engine.rs +++ b/crates/persisting-pchronicle/tests/query_engine.rs @@ -233,6 +233,12 @@ async fn atif_document_source_streams_ndjson_and_directories_in_path_order() -> .await .unwrap_err(); assert!(format!("{error:#}").contains("line 2"), "{error:#}"); + assert!( + error.chain().any(|source| source + .downcast_ref::() + .is_some()), + "missing InputIssue source: {error:#}" + ); Ok(()) } diff --git a/crates/persisting-pchronicle/tests/storyline_lance_roundtrip.rs b/crates/persisting-pchronicle/tests/storyline_lance_roundtrip.rs index 5d61588e..81b50a59 100644 --- a/crates/persisting-pchronicle/tests/storyline_lance_roundtrip.rs +++ b/crates/persisting-pchronicle/tests/storyline_lance_roundtrip.rs @@ -35,7 +35,7 @@ async fn persist_and_restore(stories: &[StorylineDocument]) -> Result Result<()> { +async fn nested_atif_and_null_canonicalization_are_stable_through_storyline_lance() -> Result<()> { let expected = serde_json::json!({ "schema_version": "ATIF-v1.7", "session_id": "shared-run", @@ -58,6 +58,7 @@ async fn nested_atif_and_null_presence_are_lossless_through_storyline_lance() -> }); let stories = decode_json_storylines(DocumentFormat::Atif, &expected.to_string(), "nested.json")?; + let expected = encode_json_storylines(DocumentFormat::Atif, &stories)?; let restored = persist_and_restore(&stories).await?; assert_eq!( encode_json_storylines(DocumentFormat::Atif, &restored)?, @@ -71,8 +72,8 @@ async fn atif_actf_and_openai_are_lossless_through_storyline_lance() -> Result<( let atif_path = Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/atif/parallel_tools_14.json"); let atif_raw = std::fs::read_to_string(&atif_path)?; - let atif_expected: serde_json::Value = serde_json::from_str(&atif_raw)?; let atif_stories = decode_json_storylines(DocumentFormat::Atif, &atif_raw, &atif_path)?; + let atif_expected = encode_json_storylines(DocumentFormat::Atif, &atif_stories)?; let atif_restored = persist_and_restore(&atif_stories).await?; assert_eq!( encode_json_storylines(DocumentFormat::Atif, &atif_restored)?, @@ -81,8 +82,8 @@ async fn atif_actf_and_openai_are_lossless_through_storyline_lance() -> Result<( let actf_path = fixture("make-doom-for-mips_trimmed.actf.json"); let actf_raw = std::fs::read_to_string(&actf_path)?; - let actf_expected: serde_json::Value = serde_json::from_str(&actf_raw)?; let actf_stories = decode_json_storylines(DocumentFormat::Actf, &actf_raw, &actf_path)?; + let actf_expected = encode_json_storylines(DocumentFormat::Actf, &actf_stories)?; let actf_restored = persist_and_restore(&actf_stories).await?; assert_eq!( encode_json_storylines(DocumentFormat::Actf, &actf_restored)?, @@ -90,11 +91,11 @@ async fn atif_actf_and_openai_are_lossless_through_storyline_lance() -> Result<( ); let openai_path = fixture("cybergym_0729001_trimmed.json"); - let openai_expected: serde_json::Value = serde_json::from_slice(&std::fs::read(&openai_path)?)?; let openai_stories = open_document(DocumentFormat::OpenaiMsg, &openai_path) .await? .project_storylines() .await?; + let openai_expected = encode_json_storylines(DocumentFormat::OpenaiMsg, &openai_stories)?; let openai_restored = persist_and_restore(&openai_stories).await?; assert_eq!( encode_json_storylines(DocumentFormat::OpenaiMsg, &openai_restored)?, @@ -149,6 +150,10 @@ async fn atif_singleton_array_shape_is_lossless_through_storyline_lance() -> Res let temporary = tempfile::tempdir()?; let input = temporary.path().join("singleton.json"); std::fs::write(&input, expected.to_string())?; + let expected = encode_json_storylines( + DocumentFormat::Atif, + &decode_json_storylines(DocumentFormat::Atif, &expected.to_string(), &input)?, + )?; let lance = temporary.path().join("storyline"); let store = StorylineLanceStore::open(&lance).await?; store.import_atif_stream(&input).await?; diff --git a/crates/persisting-pchronicle/tests/unknown_fields_roundtrip.rs b/crates/persisting-pchronicle/tests/unknown_fields_roundtrip.rs new file mode 100644 index 00000000..17c5e5ec --- /dev/null +++ b/crates/persisting-pchronicle/tests/unknown_fields_roundtrip.rs @@ -0,0 +1,58 @@ +use anyhow::Result; +use persisting_pchronicle::document::{ + decode_json_storylines, encode_json_storylines, DocumentFormat, +}; +use serde_json::{json, Value}; + +#[test] +fn atif_unknown_fields_survive_an_actf_hop() -> Result<()> { + let input = json!({ + "schema_version": "ATIF-v1.7", + "trajectory_id": "unknown-fields", + "session_id": null, + "agent": { + "name": "agent", + "version": "1", + "vendor_agent": {"enabled": true} + }, + "steps": [{ + "step_id": 1, + "source": "user", + "message": "hello", + "vendor_step": [1, {"nested": "value"}], + "vendor_null": null + }], + "vendor/root": {"tilde~key": 7} + }); + + let atif_stories = + decode_json_storylines(DocumentFormat::Atif, &input.to_string(), "source.json")?; + + let actf = encode_json_storylines(DocumentFormat::Actf, &atif_stories)?; + assert_eq!(actf["_storyline"]["unknown_fields"]["version"], 1); + + let restored = + decode_json_storylines(DocumentFormat::Actf, &actf.to_string(), "carrier.actf.json")?; + assert_eq!( + restored[0].unknown_fields.sources["atif"], + atif_stories[0].unknown_fields.sources["atif"] + ); + assert_eq!( + restored[0].unknown_key_counts["atif"]["/steps/*/vendor_step"], + 1 + ); + assert_eq!(restored[0].unknown_key_counts["atif"]["/vendor~1root"], 1); + + let recovered_atif = encode_json_storylines(DocumentFormat::Atif, &restored)?; + assert_eq!( + recovered_atif["agent"]["vendor_agent"], + json!({"enabled": true}) + ); + assert_eq!( + recovered_atif["steps"][0]["vendor_step"], + json!([1, {"nested": "value"}]) + ); + assert_eq!(recovered_atif["steps"][0]["vendor_null"], Value::Null); + assert_eq!(recovered_atif["vendor/root"], json!({"tilde~key": 7})); + Ok(()) +} diff --git a/docs/src/pchronicle/design/storyline-lance.md b/docs/src/pchronicle/design/storyline-lance.md index 921f6b7b..df8ca00e 100644 --- a/docs/src/pchronicle/design/storyline-lance.md +++ b/docs/src/pchronicle/design/storyline-lance.md @@ -358,24 +358,25 @@ SQL / DataFrame → DataFusion 保留 inexact filter 再次校验 ``` -当前 fast path 的适用范围是 ATIF 单对象 JSON、JSON 数组（包括 pretty JSON），以及每行一个对象的 JSONL/NDJSON， -目标表为 `steps`，并且物理计划存在严格列裁剪。它有意保持保守： +当前 fast path 的适用范围是 ATIF 单对象、数组（包括 pretty JSON）和 JSONL/NDJSON， +以及 ACTF 单对象和数组；目标表为 `steps`，并且物理计划存在严格列裁剪。它有意保持保守： | 输入/查询 | 执行路径 | |---|---| | ATIF object/pretty object + projected `steps` | reader-backed seeded projected decoder | | ATIF array/pretty array + projected `steps` | `fill_buf` 结构扫描 + 有界 element buffer + seeded `from_slice` | | ATIF JSONL/NDJSON + projected `steps` | `BufRead` 逐记录、有界 record buffer | +| ACTF object/array + projected `steps` | reader/slice seeded projected decoder | | `_file_`、`session_id`、`step_id`、`source` 的安全简单谓词 | 可提前裁剪，DataFusion 仍复核 | | `SELECT *` | 完整规范化 fallback | | `runs` / `tool_calls` | 完整规范化 fallback | -| OpenAI-message / ACTF | 完整规范化 fallback | +| OpenAI-message | 完整规范化 fallback | | 无法证明安全的表达式、OR/函数/跨列条件 | 不预裁剪，由 DataFusion 求值 | `DeserializeSeed` 把查询 projection 和安全谓词传入 `Visitor`；未引用字段交给 -`IgnoredAny` 做语法扫描，不构造 `Value`/Storyline。JSONL/NDJSON 以 `BufRead` 逐记录读取； -JSON array 的结构扫描器识别字符串和转义，在不构造 DOM 的情况下提取单个 trajectory， -再通过 slice decoder 执行投影解析。单条 JSONL 记录或 array element 由 +`IgnoredAny` 做语法扫描，不构造 `Value`/Storyline。ATIF JSONL/NDJSON 以 `BufRead` +逐记录读取；JSON array 的结构扫描器识别字符串和转义，在不构造 DOM 的情况下提取单个 +trajectory/document，再通过 slice decoder 执行投影解析。单条 JSONL 记录或 array element 由 `max_record_bytes` 限制；单对象直接从 reader 解码。三种路径都不先复制整文件。 Arrow encoder 也只创建投影列，`COUNT(*)` 使用合法的零列 batch。轻量路径 校验 JSON、必需字段、重复 session、命中文档内的重复 step 和当前表内约束；跨表引用 @@ -442,8 +443,8 @@ JSON 对照使用单个 NDJSON 文件，避免大量小文件打开开销。ATIF DataFusion projection 和可安全预裁剪的 `session_id`、`step_id`、`source` 谓词传给 projected decoder：未引用 JSON 字段只做语法扫描，不构造 Storyline/三表对象，Arrow batch 也只包含执行计划需要的列。object、array、pretty JSON 和 JSONL/NDJSON 共用流式 -projection decoder；`SELECT *` 和其他格式仍走完整规范化 -fallback。轻量路径执行 JSON、必需字段和表内约束校验，跨表引用完整性由导入或完整 +projection decoder；ACTF `steps` 也使用对应的 projected decoder；`SELECT *`、其他表和 +OpenAI-message 仍走完整规范化 fallback。轻量路径执行 JSON、必需字段和表内约束校验，跨表引用完整性由导入或完整 fallback 校验。预解析内存 JSON 对照只计算查询逻辑，用来区分产品工作流与纯内存遍历成本。 benchmark 还单独输出 DataSource 冷打开并执行 SQL、`get_storyline_full` 点查和单 Storyline 替换的延迟，避免 warm SQL 吞吐掩盖在线读写路径的写放大。 diff --git a/docs/src/pchronicle/design/storyline-lance.zh.md b/docs/src/pchronicle/design/storyline-lance.zh.md index 1c3114b7..45cde67b 100644 --- a/docs/src/pchronicle/design/storyline-lance.zh.md +++ b/docs/src/pchronicle/design/storyline-lance.zh.md @@ -350,24 +350,25 @@ SQL / DataFrame → DataFusion 保留 inexact filter 再次校验 ``` -当前 fast path 的适用范围是 ATIF 单对象 JSON、JSON 数组（包括 pretty JSON），以及每行一个对象的 JSONL/NDJSON， -目标表为 `steps`，并且物理计划存在严格列裁剪。它有意保持保守： +当前 fast path 的适用范围是 ATIF 单对象、数组（包括 pretty JSON）和 JSONL/NDJSON， +以及 ACTF 单对象和数组；目标表为 `steps`，并且物理计划存在严格列裁剪。它有意保持保守： | 输入/查询 | 执行路径 | |---|---| | ATIF object/pretty object + projected `steps` | reader-backed seeded projected decoder | | ATIF array/pretty array + projected `steps` | `fill_buf` 结构扫描 + 有界 element buffer + seeded `from_slice` | | ATIF JSONL/NDJSON + projected `steps` | `BufRead` 逐记录、有界 record buffer | +| ACTF object/array + projected `steps` | reader/slice seeded projected decoder | | `_file_`、`session_id`、`step_id`、`source` 的安全简单谓词 | 可提前裁剪，DataFusion 仍复核 | | `SELECT *` | 完整规范化 fallback | | `runs` / `tool_calls` | 完整规范化 fallback | -| OpenAI-message / ACTF | 完整规范化 fallback | +| OpenAI-message | 完整规范化 fallback | | 无法证明安全的表达式、OR/函数/跨列条件 | 不预裁剪，由 DataFusion 求值 | `DeserializeSeed` 把查询 projection 和安全谓词传入 `Visitor`；未引用字段交给 -`IgnoredAny` 做语法扫描，不构造 `Value`/Storyline。JSONL/NDJSON 以 `BufRead` 逐记录读取； -JSON array 的结构扫描器识别字符串和转义，在不构造 DOM 的情况下提取单个 trajectory， -再通过 slice decoder 执行投影解析。单条 JSONL 记录或 array element 由 +`IgnoredAny` 做语法扫描，不构造 `Value`/Storyline。ATIF JSONL/NDJSON 以 `BufRead` +逐记录读取；JSON array 的结构扫描器识别字符串和转义，在不构造 DOM 的情况下提取单个 +trajectory/document，再通过 slice decoder 执行投影解析。单条 JSONL 记录或 array element 由 `max_record_bytes` 限制；单对象直接从 reader 解码。三种路径都不先复制整文件。 Arrow encoder 也只创建投影列，`COUNT(*)` 使用合法的零列 batch。轻量路径 校验 JSON、必需字段、重复 session、命中文档内的重复 step 和当前表内约束；跨表引用 @@ -434,8 +435,8 @@ JSON 对照使用单个 NDJSON 文件，避免大量小文件打开开销。ATIF DataFusion projection 和可安全预裁剪的 `session_id`、`step_id`、`source` 谓词传给 projected decoder：未引用 JSON 字段只做语法扫描，不构造 Storyline/三表对象，Arrow batch 也只包含执行计划需要的列。object、array、pretty JSON 和 JSONL/NDJSON 共用流式 -projection decoder；`SELECT *` 和其他格式仍走完整规范化 -fallback。轻量路径执行 JSON、必需字段和表内约束校验，跨表引用完整性由导入或完整 +projection decoder；ACTF `steps` 也使用对应的 projected decoder；`SELECT *`、其他表和 +OpenAI-message 仍走完整规范化 fallback。轻量路径执行 JSON、必需字段和表内约束校验，跨表引用完整性由导入或完整 fallback 校验。预解析内存 JSON 对照只计算查询逻辑，用来区分产品工作流与纯内存遍历成本。 benchmark 还单独输出 DataSource 冷打开并执行 SQL、`get_storyline_full` 点查和单 Storyline 替换的延迟，避免 warm SQL 吞吐掩盖在线读写路径的写放大。 diff --git a/docs/src/pchronicle/design/trajectory-storage.md b/docs/src/pchronicle/design/trajectory-storage.md index 2cbdbc43..8449dd72 100644 --- a/docs/src/pchronicle/design/trajectory-storage.md +++ b/docs/src/pchronicle/design/trajectory-storage.md @@ -80,11 +80,11 @@ result 归并到 tool call 行，并通过 `CURRENT` 中的三表版本元组保 UTF-8/JSON cell 以 BLAKE3 内容地址外置到共享 `objects.lance`，跨轨迹复用；公开 schema 和 SQL 结果保持不变，查询只在真正引用内容列时延迟恢复 Blob。 -ATIF object、array、pretty JSON 与 JSONL/NDJSON 的 `steps` 临时查询还支持 +ATIF object、array、pretty JSON 与 JSONL/NDJSON，以及 ACTF object/array 的 `steps` 临时查询还支持 projection-aware 快路径：DataFusion 先传递所需列和安全谓词，reader 通过 seeded visitor 跳过未引用大字段并直接构造窄 Arrow batch；JSONL/NDJSON 逐记录有界读取，array 通过 结构扫描器逐 element 提取并使用 slice decoder，单 object 从 reader 流式解码。 -`SELECT *`、其他表和格式回退到完整 Storyline 规范化。详细协议、发布顺序和执行边界见 +`SELECT *`、其他表和 OpenAI-message 回退到完整 Storyline 规范化。详细协议、发布顺序和执行边界见 [Storyline 三表 Lance 存储](storyline-lance.md)。 ## 4. 目录布局 diff --git a/docs/src/pchronicle/design/trajectory-storage.zh.md b/docs/src/pchronicle/design/trajectory-storage.zh.md index 2cbdbc43..8449dd72 100644 --- a/docs/src/pchronicle/design/trajectory-storage.zh.md +++ b/docs/src/pchronicle/design/trajectory-storage.zh.md @@ -80,11 +80,11 @@ result 归并到 tool call 行，并通过 `CURRENT` 中的三表版本元组保 UTF-8/JSON cell 以 BLAKE3 内容地址外置到共享 `objects.lance`，跨轨迹复用；公开 schema 和 SQL 结果保持不变，查询只在真正引用内容列时延迟恢复 Blob。 -ATIF object、array、pretty JSON 与 JSONL/NDJSON 的 `steps` 临时查询还支持 +ATIF object、array、pretty JSON 与 JSONL/NDJSON，以及 ACTF object/array 的 `steps` 临时查询还支持 projection-aware 快路径：DataFusion 先传递所需列和安全谓词，reader 通过 seeded visitor 跳过未引用大字段并直接构造窄 Arrow batch；JSONL/NDJSON 逐记录有界读取，array 通过 结构扫描器逐 element 提取并使用 slice decoder，单 object 从 reader 流式解码。 -`SELECT *`、其他表和格式回退到完整 Storyline 规范化。详细协议、发布顺序和执行边界见 +`SELECT *`、其他表和 OpenAI-message 回退到完整 Storyline 规范化。详细协议、发布顺序和执行边界见 [Storyline 三表 Lance 存储](storyline-lance.md)。 ## 4. 目录布局 diff --git a/docs/src/rfcs/0004-actf-format.md b/docs/src/rfcs/0004-actf-format.md index 09f11558..8b4bbddf 100644 --- a/docs/src/rfcs/0004-actf-format.md +++ b/docs/src/rfcs/0004-actf-format.md @@ -61,7 +61,8 @@ ActfObservation = { type: string, id?: string, tool_use_id?: string, ``` token 数、`llm_infer_ms` 和 `env_action_ms` 可以是数值或 `null`，`stop_reason` 也可以 -显式为 `null`；实现不得把显式 `null` 静默改成字段缺失。ACTF v1.0 已观察到两种工具事件： +显式为 `null`；解析器必须接受两种表示，进入 Storyline 后 missing/null 按同一语义默认值 +规范化。ACTF v1.0 已观察到两种工具事件： `tool_use` 使用 `name/input` 与 `tool_use_id/content`，`command_execution` 使用 `command/aggregated_output/exit_code/status`，并通过共同的 `id` 关联。事件专属字段作为 opaque JSON 保留。 @@ -93,13 +94,15 @@ ACTF tool 1 ──► 1 Storyline tool_call `{task_id}#attempt-{attempt_id}`。 - assistant `content`、`reasoning_content`、tool call 和 observation 投影到对应的 Storyline 字段，token/latency 投影到 metrics。 -- ACTF 根、attempt、trajectory 元数据和完整原始 step 写入三表现有的 `extra_json` - 扩展列，并带 `_pchronicle_actf` provenance version。 -- 恢复时以 provenance 重组 attempt map；多个 attempt 必须具有相同根元数据。 +- Storyline 未建模的 ACTF 根、attempt、trajectory 和 step 键以精确 RFC 6901 JSON + Pointer/value 写入 run 级 `unknown_fields`；`unknown_key_counts` 记录归一化路径的出现次数。 +- 恢复时根据 `run_id`、`attempt_id` 和 ACTF unknown fields 重组 attempt map；多个 attempt + 必须具有相同根元数据。跨格式转换通过 version-1 `_storyline` envelope 携带 unknown fields。 ## 保真边界 -ACTF → Storyline → 三表 Lance → Storyline → ACTF 保证 JSON 数据模型级无损：键值、显式 -`null`、未知字段、嵌套值、数组顺序和 attempt 分组均保留。它不保证源文件空白、缩进或 -对象键顺序逐字节一致。没有 ACTF provenance 的普通 Storyline 可以导出为结构合法的 -单 attempt ACTF，但这是有定义的合成转换，不宣称还原某个原始 ACTF 文件。 +ACTF → Storyline → 三表 Lance → Storyline → ACTF 保证规范化 JSON 数据模型级语义一致： +未知键及其值（包括 `null`）、嵌套值、数组顺序和 attempt 分组均保留。已知字段的 +missing/显式 `null` 会按 Storyline 语义规范化；源文件空白、缩进和对象键顺序不属于保真 +边界。没有 ACTF source unknown fields 的普通 Storyline 可以导出为结构合法的单 attempt ACTF， +但这是有定义的合成转换，不宣称还原某个原始 ACTF 文件。 From 6b8fa48065d5c8399a708909726e45f714d3b7ab Mon Sep 17 00:00:00 2001 From: Reiase Date: Thu, 20 Aug 2026 10:02:04 +0800 Subject: [PATCH 2/7] refactor(pchronicle): close review gaps --- crates/persisting-dlcapt/src/tlv.rs | 24 ++++ .../src/agenticmd/convert.rs | 8 +- .../persisting-pchronicle/src/convert/actf.rs | 18 +-- .../persisting-pchronicle/src/convert/atif.rs | 1 + crates/persisting-pchronicle/src/document.rs | 42 +++++-- .../src/formats/openai_corpus.rs | 15 +-- .../src/formats/unknown_fields.rs | 24 +++- .../src/store/files/mod.rs | 5 +- .../src/store/storyline/mod.rs | 10 ++ .../src/store/storyline/tests.rs | 38 +++++- .../tests/conversion_semantics.rs | 115 +++++++++++++++++- docs/src/pchronicle/design/storyline-lance.md | 7 +- .../pchronicle/design/storyline-lance.zh.md | 7 +- .../pchronicle/design/trajectory-storage.md | 3 +- .../design/trajectory-storage.zh.md | 3 +- 15 files changed, 273 insertions(+), 47 deletions(-) diff --git a/crates/persisting-dlcapt/src/tlv.rs b/crates/persisting-dlcapt/src/tlv.rs index 684beb79..524fb36d 100644 --- a/crates/persisting-dlcapt/src/tlv.rs +++ b/crates/persisting-dlcapt/src/tlv.rs @@ -218,6 +218,13 @@ fn encode_block( kind: &str, ) -> Result { let timestamp = Utc::now().to_rfc3339(); + let storyline_id = + i64::try_from(seq).context("tlv sequence exceeds Storyline turn id range")?; + let storyline_source = if speaker == "assistant" { + "agent" + } else { + speaker + }; let mut fields = BTreeMap::new(); fields.insert("agent_id".to_string(), json!(record.agent_id)); fields.insert("call_id".to_string(), json!(record.call_id)); @@ -232,6 +239,18 @@ fn encode_block( fields.insert("trace_id".to_string(), json!(record.call_id)); fields.insert("turn".to_string(), json!(record.turn)); fields.insert("v".to_string(), json!(BLOCK_FORMAT_VERSION)); + fields.insert("message_encoding".to_string(), json!("text")); + fields.insert("step_id".to_string(), json!(storyline_id)); + fields.insert( + "storyline".to_string(), + json!({ + "id": storyline_id, + "kind": kind, + "ts": timestamp, + "src": storyline_source, + "model": record.model, + }), + ); if speaker == "assistant" { fields.insert("status".to_string(), json!(record.status_code)); @@ -290,6 +309,11 @@ fn format_document_preamble(session_id: &str, agent_id: &str, turns: u64) -> Res "session: {session_id}\n", "agent: {agent_id}\n", "turns: {turns}\n", + "storyline:\n", + " session: {session_id}\n", + " agent:\n", + " id: {agent_id}\n", + " name: {agent_id}\n", "client:\n", " peer: ''\n", " peer_port: 0\n", diff --git a/crates/persisting-pchronicle/src/agenticmd/convert.rs b/crates/persisting-pchronicle/src/agenticmd/convert.rs index 720c42e6..2680dc94 100644 --- a/crates/persisting-pchronicle/src/agenticmd/convert.rs +++ b/crates/persisting-pchronicle/src/agenticmd/convert.rs @@ -212,11 +212,13 @@ fn capture_agenticmd_unknown_fields( fn agenticmd_source_document_id(document: &MarkdownDocument) -> InputResult { let mut source = serde_json::to_value(document).map_err(|error| InputIssue::invalid(error.to_string()))?; - source + let frontmatter = source .get_mut("frontmatter") .and_then(Value::as_object_mut) - .expect("serialized AgenticMD document has object frontmatter") - .remove(STORYLINE_METADATA_KEY); + .ok_or_else(|| { + InputIssue::invalid("serialized AgenticMD document lacks object frontmatter") + })?; + frontmatter.remove(STORYLINE_METADATA_KEY); canonical_source_document_id(&source).map_err(|error| InputIssue::invalid(error.to_string())) } diff --git a/crates/persisting-pchronicle/src/convert/actf.rs b/crates/persisting-pchronicle/src/convert/actf.rs index 2283d00d..3704dcdb 100644 --- a/crates/persisting-pchronicle/src/convert/actf.rs +++ b/crates/persisting-pchronicle/src/convert/actf.rs @@ -162,7 +162,7 @@ fn capture_actf_unknowns( .map_err(|error| crate::InputIssue::invalid(error.to_string()))?; let root = root .as_object_mut() - .expect("ACTF root serializes as object"); + .ok_or_else(|| crate::InputIssue::invalid("serialized ACTF document must be an object"))?; for key in ["task_id", "correct", "attempts"] { root.remove(key); } @@ -173,7 +173,7 @@ fn capture_actf_unknowns( .map_err(|error| crate::InputIssue::invalid(error.to_string()))?; let attempt_map = attempt_value .as_object_mut() - .expect("ACTF attempt is object"); + .ok_or_else(|| crate::InputIssue::invalid("serialized ACTF attempt must be an object"))?; for key in ["correct", "score", "status", "trajectory"] { attempt_map.remove(key); } @@ -182,9 +182,9 @@ fn capture_actf_unknowns( let trajectory_prefix = pointer_join(&attempt_prefix, "trajectory"); let mut trajectory_value = serde_json::to_value(&attempt.trajectory) .map_err(|error| crate::InputIssue::invalid(error.to_string()))?; - let trajectory_map = trajectory_value - .as_object_mut() - .expect("ACTF trajectory is object"); + let trajectory_map = trajectory_value.as_object_mut().ok_or_else(|| { + crate::InputIssue::invalid("serialized ACTF trajectory must be an object") + })?; trajectory_map.remove("steps"); insert_actf_map(story, source_id, &trajectory_prefix, trajectory_map)?; @@ -195,14 +195,18 @@ fn capture_actf_unknowns( ); let mut step_value = serde_json::to_value(step) .map_err(|error| crate::InputIssue::invalid(error.to_string()))?; - let step_map = step_value.as_object_mut().expect("ACTF step is object"); + let step_map = step_value + .as_object_mut() + .ok_or_else(|| crate::InputIssue::invalid("serialized ACTF step must be an object"))?; let assistant = step_map.remove("assistant_content"); for key in ["step_id", "metric", "tools", "observation", "started_at"] { step_map.remove(key); } insert_actf_map(story, source_id, &step_prefix, step_map)?; if let Some(mut assistant) = assistant { - let assistant = assistant.as_object_mut().expect("ACTF assistant is object"); + let assistant = assistant.as_object_mut().ok_or_else(|| { + crate::InputIssue::invalid("serialized ACTF assistant content must be an object") + })?; for key in ["content", "reasoning_content", "tool_calls"] { assistant.remove(key); } diff --git a/crates/persisting-pchronicle/src/convert/atif.rs b/crates/persisting-pchronicle/src/convert/atif.rs index 0b4a0cb9..6f912d89 100644 --- a/crates/persisting-pchronicle/src/convert/atif.rs +++ b/crates/persisting-pchronicle/src/convert/atif.rs @@ -127,6 +127,7 @@ pub(crate) fn atif_value_to_storylines( let (mut stories, carriers) = atif_to_storylines_with_source(&trajectory, &source_document_id) .map_err(|error| crate::InputIssue::invalid(error.to_string()))?; attach_carried_unknown_fields( + DocumentFormat::Atif, envelope, &carriers, &mut stories, diff --git a/crates/persisting-pchronicle/src/document.rs b/crates/persisting-pchronicle/src/document.rs index 1f9d734c..a5372088 100644 --- a/crates/persisting-pchronicle/src/document.rs +++ b/crates/persisting-pchronicle/src/document.rs @@ -92,10 +92,6 @@ pub fn decode_json_storylines_with_options( document.validate()?; let mut stories = crate::convert::actf_to_storylines(&document) .map_err(|error| InputIssue::invalid(error.to_string()))?; - let owned_counts = stories - .iter() - .map(|story| story.unknown_key_counts.get("actf").cloned()) - .collect::>(); let carriers = stories .iter() .enumerate() @@ -113,16 +109,12 @@ pub fn decode_json_storylines_with_options( }) .collect::>(); attach_carried_unknown_fields( + DocumentFormat::Actf, envelope, &carriers, &mut stories, options.unknown_fields, )?; - for (story, owned) in stories.iter_mut().zip(owned_counts) { - if let Some(owned) = owned { - story.unknown_key_counts.insert("actf".into(), owned); - } - } Ok(stories) } DocumentFormat::OpenaiMsg => { @@ -340,6 +332,38 @@ mod tests { .is_err()); } + #[test] + fn openai_envelope_rejects_target_source_unknown_fields() { + let input = serde_json::json!({ + "session_steps": [{ + "session_id": "s-1", + "step_id": 1, + "messages": [{"role": "user", "content": "inspect"}], + "response": {"role": "assistant", "content": "done"} + }], + "_storyline": {"unknown_fields": { + "version": 1, + "by_trajectory": {"/session_steps/0": {"sources": { + "openai-msg": { + "source_document_id": "same-target.json", + "fields": {"/session_steps/0/enveloped_vendor": true} + } + }}} + }} + }); + + let error = decode_json_storylines( + DocumentFormat::OpenaiMsg, + &input.to_string(), + "same-target.json", + ) + .unwrap_err(); + assert!( + error.to_string().contains("target source 'openai-msg'"), + "{error:#}" + ); + } + #[test] fn unknown_fields_options_are_validated_before_json_decode() { let options = DocumentCodecOptions { diff --git a/crates/persisting-pchronicle/src/formats/openai_corpus.rs b/crates/persisting-pchronicle/src/formats/openai_corpus.rs index 2a03b1d7..a9501bb8 100644 --- a/crates/persisting-pchronicle/src/formats/openai_corpus.rs +++ b/crates/persisting-pchronicle/src/formats/openai_corpus.rs @@ -114,21 +114,13 @@ pub fn parse_openai_msg_corpus_value( } stories.push(story); } - let owned_counts = stories - .iter() - .map(|story| story.unknown_key_counts.get("openai-msg").cloned()) - .collect::>(); attach_carried_unknown_fields( + DocumentFormat::OpenaiMsg, carried_envelope, &carriers, &mut stories, UnknownFieldLimits::default(), )?; - for (story, owned) in stories.iter_mut().zip(owned_counts) { - if let Some(owned) = owned { - story.unknown_key_counts.insert("openai-msg".into(), owned); - } - } Ok(stories) } @@ -148,7 +140,10 @@ fn capture_openai_unknowns( }); insert_openai_map(story, source_document_id, "", root_unknown)?; for (ordinal, record) in records { - let row = record.as_object().expect("rows were validated as objects"); + let row = record.as_object().ok_or_else(|| { + InputIssue::invalid("OpenAI corpus row must be an object") + .at(format!("rows[{ordinal}]")) + })?; let row_prefix = format!("/session_steps/{ordinal}"); for (key, value) in row { if !is_canonical_openai_row_key(key) { diff --git a/crates/persisting-pchronicle/src/formats/unknown_fields.rs b/crates/persisting-pchronicle/src/formats/unknown_fields.rs index 805516b4..b0e63e94 100644 --- a/crates/persisting-pchronicle/src/formats/unknown_fields.rs +++ b/crates/persisting-pchronicle/src/formats/unknown_fields.rs @@ -127,14 +127,15 @@ pub(crate) fn take_unknown_fields_envelope( carried.insert(carrier, fields); } - document + let document = document .as_object_mut() - .expect("an object containing _storyline remains an object") - .remove("_storyline"); + .ok_or_else(|| InputIssue::invalid("_storyline envelope requires an object document"))?; + document.remove("_storyline"); Ok(carried) } pub(crate) fn attach_carried_unknown_fields( + target_format: DocumentFormat, envelope: BTreeMap, carriers: &[CarrierBinding], stories: &mut [StorylineDocument], @@ -149,6 +150,12 @@ pub(crate) fn attach_carried_unknown_fields( .collect::>(); for (pointer, carried) in envelope { validate_json_pointer(&pointer)?; + let target_source = target_format.as_str(); + if carried.sources.contains_key(target_source) { + return Err(InputIssue::invalid(format!( + "_storyline envelope must not carry target source '{target_source}'" + ))); + } let story_index = by_pointer.get(&pointer).ok_or_else(|| { InputIssue::invalid(format!( "_storyline envelope carrier '{pointer}' is not bound to a trajectory" @@ -658,9 +665,9 @@ pub(crate) fn restore_json_pointer( }; } - let (last, parents) = tokens - .split_last() - .expect("non-empty JSON Pointer token sequence"); + let Some((last, parents)) = tokens.split_last() else { + anyhow::bail!("JSON Pointer unexpectedly decoded to an empty token sequence"); + }; let mut parent = target; for token in parents { parent = match parent { @@ -762,6 +769,7 @@ mod tests { assert!(raw.get("_storyline").is_none()); let mut stories = vec![StorylineDocument::new("s", "a")]; attach_carried_unknown_fields( + DocumentFormat::Actf, envelope, &[CarrierBinding { story_index: 0, @@ -810,6 +818,7 @@ mod tests { let stories = &mut [StorylineDocument::new("s", "a")]; assert!(attach_carried_unknown_fields( + DocumentFormat::OpenaiMsg, BTreeMap::new(), &[ CarrierBinding { @@ -828,6 +837,7 @@ mod tests { let unbound = BTreeMap::from([("/missing".into(), StorylineUnknownFields::default())]); assert!(attach_carried_unknown_fields( + DocumentFormat::OpenaiMsg, unbound, &[CarrierBinding { story_index: 0, @@ -861,6 +871,7 @@ mod tests { }, )]); assert!(attach_carried_unknown_fields( + DocumentFormat::OpenaiMsg, changed_id, &[CarrierBinding { story_index: 0, @@ -888,6 +899,7 @@ mod tests { }, )]); assert!(attach_carried_unknown_fields( + DocumentFormat::OpenaiMsg, carried, &[CarrierBinding { story_index: 0, diff --git a/crates/persisting-pchronicle/src/store/files/mod.rs b/crates/persisting-pchronicle/src/store/files/mod.rs index da3e8b37..bd9a4d27 100644 --- a/crates/persisting-pchronicle/src/store/files/mod.rs +++ b/crates/persisting-pchronicle/src/store/files/mod.rs @@ -824,7 +824,10 @@ fn load_file( runtime.options.max_record_bytes, ) .with_context(|| format!("parse ACTF input {}", state.file.path().display()))?, - _ => unreachable!(), + _ => anyhow::bail!( + "{} does not support streaming trajectory parsing", + format.as_str() + ), }; runtime.metrics.inner.streaming_buffer_peak_bytes.fetch_max( (reader.capacity() as u64).saturating_add(peak_record_bytes as u64), diff --git a/crates/persisting-pchronicle/src/store/storyline/mod.rs b/crates/persisting-pchronicle/src/store/storyline/mod.rs index 95d1aa09..a6e97717 100644 --- a/crates/persisting-pchronicle/src/store/storyline/mod.rs +++ b/crates/persisting-pchronicle/src/store/storyline/mod.rs @@ -87,6 +87,7 @@ use super::{root_write_lock, LanceMaintenanceOptions, LanceMaintenanceReport}; const CURRENT_FILE: &str = "CURRENT"; const GENERATIONS_DIR: &str = "generations"; +const STORYLINE_LANCE_SCHEMA_VERSION: u32 = 1; const WRITE_BATCH_ROWS: usize = 8192; const STREAM_IMPORT_STORIES: usize = 256; const RUN_INDEXES: [(&str, IndexType); 3] = [ @@ -191,6 +192,7 @@ impl StorylineProjectionLineage { #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] struct StorylineSnapshotPointer { + schema_version: u32, generation: String, #[serde(default, skip_serializing_if = "Option::is_none")] parent_generation: Option, @@ -547,6 +549,12 @@ impl StorylineLanceStore { } let pointer = serde_json::from_str::(contents) .context("decode Storyline snapshot pointer")?; + anyhow::ensure!( + pointer.schema_version == STORYLINE_LANCE_SCHEMA_VERSION, + "unsupported Storyline Lance schema_version {}; expected {}", + pointer.schema_version, + STORYLINE_LANCE_SCHEMA_VERSION + ); if let Some(projection) = &pointer.projection { projection.validate()?; } @@ -905,6 +913,7 @@ impl StorylineLanceStore { }; let generation = next_generation(); let snapshot = StorylineSnapshotPointer { + schema_version: STORYLINE_LANCE_SCHEMA_VERSION, generation: generation.clone(), parent_generation: expected_generation.clone(), table_generation: current.table_generation.clone(), @@ -1004,6 +1013,7 @@ impl StorylineLanceStore { let generation = next_generation(); self.commit_snapshot( &StorylineSnapshotPointer { + schema_version: STORYLINE_LANCE_SCHEMA_VERSION, generation: generation.clone(), parent_generation: Some(paths.generation.clone()), table_generation: paths.table_generation.clone(), diff --git a/crates/persisting-pchronicle/src/store/storyline/tests.rs b/crates/persisting-pchronicle/src/store/storyline/tests.rs index 75bf776e..48ae679d 100644 --- a/crates/persisting-pchronicle/src/store/storyline/tests.rs +++ b/crates/persisting-pchronicle/src/store/storyline/tests.rs @@ -1346,8 +1346,8 @@ async fn open_rejects_malformed_or_incomplete_commit_pointer() { tokio::fs::write( complete_pointer.path().join(CURRENT_FILE), serde_json::to_vec(&serde_json::json!({ - "generation": "snapshot-current", - "table_generation": "tables-current", + "generation": "gen-1-1-1", + "table_generation": "gen-1-1-1", "runs_version": 1, "steps_version": 1, "tool_calls_version": 1, @@ -1363,6 +1363,39 @@ async fn open_rejects_malformed_or_incomplete_commit_pointer() { assert!(!error.to_string().is_empty()); } +#[tokio::test] +async fn open_rejects_missing_or_unsupported_snapshot_schema_version() { + for (schema_version, expected) in [ + (None, "schema_version"), + ( + Some(2), + "unsupported Storyline Lance schema_version 2; expected 1", + ), + ] { + let root = tempfile::tempdir().unwrap(); + let mut pointer = serde_json::json!({ + "generation": "gen-1-1-1", + "table_generation": "gen-1-1-1", + "runs_version": 1, + "steps_version": 1, + "tool_calls_version": 1, + "objects_version": 1 + }); + if let Some(schema_version) = schema_version { + pointer["schema_version"] = serde_json::json!(schema_version); + } + tokio::fs::write( + root.path().join(CURRENT_FILE), + serde_json::to_vec(&pointer).unwrap(), + ) + .await + .unwrap(); + + let error = StorylineLanceStore::open(root.path()).await.unwrap_err(); + assert!(format!("{error:#}").contains(expected), "{error:#}"); + } +} + #[tokio::test] async fn object_store_uri_round_trips_across_store_instances() { let uri = format!("{}/", remote_uri("round-trip")); @@ -1661,6 +1694,7 @@ async fn stale_current_commit_is_rejected_without_moving_snapshot() { let error = store .commit_snapshot( &StorylineSnapshotPointer { + schema_version: STORYLINE_LANCE_SCHEMA_VERSION, generation: attempted_generation, parent_generation: Some(stale.generation.clone()), table_generation: stale.table_generation.clone(), diff --git a/crates/persisting-pchronicle/tests/conversion_semantics.rs b/crates/persisting-pchronicle/tests/conversion_semantics.rs index 0a4cc9df..a0bc16bb 100644 --- a/crates/persisting-pchronicle/tests/conversion_semantics.rs +++ b/crates/persisting-pchronicle/tests/conversion_semantics.rs @@ -317,6 +317,118 @@ fn common_storyline_semantics(stories: &[StorylineDocument]) -> Value { ) } +fn assert_target_modeled_semantics( + source: &str, + target: DocumentFormat, + stories: &[StorylineDocument], +) { + assert_eq!(stories.len(), 1); + let story = &stories[0]; + match (source, target) { + ("atif", DocumentFormat::Actf) => { + assert_eq!(story.agent.id, "actf-agent"); + assert_eq!(story.agent.name.as_deref(), Some("ACTF Agent")); + assert!(story.turns.iter().all(|turn| turn.source == "agent")); + assert!(story + .turns + .iter() + .all(|turn| turn.timestamp.as_deref() == Some("1970-01-01 00:00:00+00:00"))); + assert!(story + .turns + .iter() + .all(|turn| turn.reasoning_content.is_none())); + assert_eq!( + story.turns[1].metrics.as_ref().unwrap()["prompt_tokens_len"], + 0 + ); + assert_eq!( + story.turns[1].metrics.as_ref().unwrap()["completion_tokens_len"], + 0 + ); + } + ("atif", DocumentFormat::OpenaiMsg) => { + assert_eq!(story.agent.id, "agent"); + assert_eq!(story.agent.name.as_deref(), Some("agent")); + assert!(story.agent.version.is_none()); + assert_eq!(story.turns[0].source, "user"); + assert_eq!(story.turns[1].source, "agent"); + assert!(story.turns.iter().all(|turn| turn.timestamp.is_none())); + assert_eq!(story.turns[1].metrics.as_ref().unwrap()["reward"], 0.0); + } + ("actf", DocumentFormat::Atif) => { + assert_eq!(story.agent.id, "ACTF Agent"); + assert_eq!(story.agent.name.as_deref(), Some("ACTF Agent")); + assert_eq!(story.agent.version.as_deref(), Some("unknown")); + assert_eq!(story.turns[0].reasoning_content.as_deref(), Some("inspect")); + assert_eq!( + story.turns[1].reasoning_content.as_deref(), + Some("complete") + ); + assert_eq!( + story.turns[1].timestamp.as_deref(), + Some("2026-08-20T00:00:01Z") + ); + assert_eq!( + story.turns[1].metrics.as_ref().unwrap()["prompt_tokens_len"], + 10 + ); + assert_eq!( + story.turns[1].metrics.as_ref().unwrap()["completion_tokens_len"], + 4 + ); + assert_eq!(story.turns[1].metrics.as_ref().unwrap()["llm_infer_ms"], 12); + assert_eq!(story.turns[1].latency_ms, Some(12)); + assert_eq!( + story.turns[1].tool_calls.as_ref().unwrap()[0].duration_ms, + Some(3) + ); + } + ("actf", DocumentFormat::OpenaiMsg) => { + assert_eq!(story.agent.id, "actf-agent"); + assert_eq!(story.agent.name.as_deref(), Some("actf-agent")); + assert!(story + .turns + .iter() + .all(|turn| turn.reasoning_content.is_none())); + assert_eq!( + story.turns[1].timestamp.as_deref(), + Some("2026-08-20T00:00:01Z") + ); + assert_eq!(story.turns[1].metrics.as_ref().unwrap()["reward"], 0.0); + assert_eq!( + story.turns[1].tool_calls.as_ref().unwrap()[0].duration_ms, + None + ); + } + ("openai", DocumentFormat::Atif) => { + assert_eq!(story.agent.id, "model"); + assert_eq!(story.agent.model_name.as_deref(), Some("model")); + assert_eq!(story.agent.version.as_deref(), Some("unknown")); + assert!(story + .turns + .iter() + .all(|turn| turn.timestamp.as_deref() == Some("2026-08-20T00:00:00Z"))); + assert_eq!(story.turns[1].model_name.as_deref(), Some("model")); + assert!(story.turns[1].metrics.as_ref().unwrap()["reward"].is_null()); + } + ("openai", DocumentFormat::Actf) => { + assert_eq!(story.agent.id, "actf-agent"); + assert_eq!(story.agent.name.as_deref(), Some("ACTF Agent")); + assert!(story.turns.iter().all(|turn| turn.source == "agent")); + assert!(story + .turns + .iter() + .all(|turn| turn.timestamp.as_deref() == Some("2026-08-20 00:00:00+00:00"))); + assert_eq!( + story.turns[1].metrics.as_ref().unwrap()["prompt_tokens_len"], + 0 + ); + assert!(story.turns[1].metrics.as_ref().unwrap()["reward"].is_null()); + } + _ => panic!("unexpected conversion {source} -> {target}"), + } +} + fn assert_common_storyline_semantics(actual: &[StorylineDocument], expected: &[StorylineDocument]) { assert_semantic_json_eq( &common_storyline_semantics(actual), @@ -405,13 +517,14 @@ async fn every_directed_cross_format_hop_is_semantically_stable_through_lance() .with_context(|| format!("decode {edge}"))?; assert_unknown_contract(source.format, &bridge_stories); assert_common_storyline_semantics(&bridge_stories, &source_stories); + assert_target_modeled_semantics(source.name, bridge, &bridge_stories); let restored = persist_and_restore(&bridge_stories) .await .with_context(|| format!("persist {edge}"))?; assert_eq!(restored, bridge_stories, "Storyline changed across {edge}"); assert_unknown_contract(source.format, &restored); - assert_common_storyline_semantics(&restored, &source_stories); + assert_target_modeled_semantics(source.name, bridge, &restored); let actual = encode_json_storylines(bridge, &restored) .with_context(|| format!("restore {edge} bridge"))?; assert_semantic_json_eq(&actual, &bridged); diff --git a/docs/src/pchronicle/design/storyline-lance.md b/docs/src/pchronicle/design/storyline-lance.md index df8ca00e..d36bf4e4 100644 --- a/docs/src/pchronicle/design/storyline-lance.md +++ b/docs/src/pchronicle/design/storyline-lance.md @@ -198,8 +198,8 @@ root/ `replace_storyline` 不再读取或重写全库，而是按各表主键执行 merge-upsert，并只删除指定 `document_id` 中已经不再存在的旧键。每次替换 会产生一个新的逻辑 snapshot； -`CURRENT` 是一段 JSON，记录逻辑 snapshot id、物理 `table_generation`、三张表以及对象表 -各自精确的 Lance version id。对象先持久化，三张业务表随后写入，最后才更新 `CURRENT`； +`CURRENT` 是一段 JSON，记录必需的 store `schema_version: 1`、逻辑 snapshot id、物理 +`table_generation`、三张表以及对象表各自精确的 Lance version id。对象先持久化，三张业务表随后写入，最后才更新 `CURRENT`； 因此失败最多留下不可达对象，不会发布悬空引用或跨表半提交。 阈值、preview 长度和 Zstd level 可通过 `StorylineContentOptions` 配置；三表 schema 不变。 @@ -209,7 +209,8 @@ Lance MVCC 的旧版本默认保留，便于已打开的 reader 固定快照及 compaction，避免某次写请求出现维护型长尾；生产环境通过 `maintain` 显式执行三表并行 compaction、补齐/刷新索引、内容 GC 和按保留期 vacuum。维护产生的四个 dataset version 仍先原子更新 `CURRENT`，之后才回收旧版本和过期的非当前 physical generation。 -`CURRENT` 必须是包含全部精确版本的 JSON 指针，不读取旧的纯文本 generation 指针。 +`CURRENT` 必须是包含 schema version 和全部精确版本的 JSON 指针；缺失或未知 schema +version 会在打开任何 Lance table 前 fail closed，也不读取旧的纯文本 generation 指针。 本地写入通过进程内锁和文件锁串行化；对象存储通过 `CURRENT` 的 ETag/version 条件更新 执行 optimistic CAS。stale commit 不能移动 `CURRENT`；`StorylineLanceStore` 在 CAS 冲突后 diff --git a/docs/src/pchronicle/design/storyline-lance.zh.md b/docs/src/pchronicle/design/storyline-lance.zh.md index 45cde67b..3fd1906b 100644 --- a/docs/src/pchronicle/design/storyline-lance.zh.md +++ b/docs/src/pchronicle/design/storyline-lance.zh.md @@ -196,8 +196,8 @@ root/ `replace_storyline` 不再读取或重写全库，而是按各表主键执行 merge-upsert，并只删除指定 `document_id` 中已经不再存在的旧键。每次替换 会产生一个新的逻辑 snapshot； -`CURRENT` 是一段 JSON，记录逻辑 snapshot id、物理 `table_generation`、三张表以及对象表 -各自精确的 Lance version id。对象先持久化，三张业务表随后写入，最后才更新 `CURRENT`； +`CURRENT` 是一段 JSON，记录必需的 store `schema_version: 1`、逻辑 snapshot id、物理 +`table_generation`、三张表以及对象表各自精确的 Lance version id。对象先持久化，三张业务表随后写入，最后才更新 `CURRENT`； 因此失败最多留下不可达对象，不会发布悬空引用或跨表半提交。 阈值、preview 长度和 Zstd level 可通过 `StorylineContentOptions` 配置；三表 schema 不变。 @@ -207,7 +207,8 @@ Lance MVCC 的旧版本默认保留，便于已打开的 reader 固定快照及 compaction，避免某次写请求出现维护型长尾；生产环境通过 `maintain` 显式执行三表并行 compaction、补齐/刷新索引、内容 GC 和按保留期 vacuum。维护产生的四个 dataset version 仍先原子更新 `CURRENT`，之后才回收旧版本和过期的非当前 physical generation。 -`CURRENT` 必须是包含全部精确版本的 JSON 指针，不读取旧的纯文本 generation 指针。 +`CURRENT` 必须是包含 schema version 和全部精确版本的 JSON 指针；缺失或未知 schema +version 会在打开任何 Lance table 前 fail closed，也不读取旧的纯文本 generation 指针。 本地写入通过进程内锁和文件锁串行化；对象存储通过 `CURRENT` 的 ETag/version 条件更新 执行 optimistic CAS。stale commit 不能移动 `CURRENT`；`StorylineLanceStore` 在 CAS 冲突后 diff --git a/docs/src/pchronicle/design/trajectory-storage.md b/docs/src/pchronicle/design/trajectory-storage.md index 8449dd72..698375e7 100644 --- a/docs/src/pchronicle/design/trajectory-storage.md +++ b/docs/src/pchronicle/design/trajectory-storage.md @@ -113,7 +113,8 @@ storage/ ``` 独立的 Storyline 分析 store 使用 `CURRENT`、`generations//{runs,steps,tool_calls}.lance` -和根级共享 `objects.lance`；它不改变上面的 canonical event 目录。 +和根级共享 `objects.lance`；`CURRENT` 的必需 `schema_version` 在打开表前验证整个四表物理 +布局。它不改变上面的 canonical event 目录。 系统生成的 AgenticMD 使用 `{session_id}.md` 文件名和 `` 块结构；读取器同时接受无 speaker 的块、旧 diff --git a/docs/src/pchronicle/design/trajectory-storage.zh.md b/docs/src/pchronicle/design/trajectory-storage.zh.md index 8449dd72..698375e7 100644 --- a/docs/src/pchronicle/design/trajectory-storage.zh.md +++ b/docs/src/pchronicle/design/trajectory-storage.zh.md @@ -113,7 +113,8 @@ storage/ ``` 独立的 Storyline 分析 store 使用 `CURRENT`、`generations//{runs,steps,tool_calls}.lance` -和根级共享 `objects.lance`；它不改变上面的 canonical event 目录。 +和根级共享 `objects.lance`；`CURRENT` 的必需 `schema_version` 在打开表前验证整个四表物理 +布局。它不改变上面的 canonical event 目录。 系统生成的 AgenticMD 使用 `{session_id}.md` 文件名和 `` 块结构；读取器同时接受无 speaker 的块、旧 From 455b77c14fc582ab9a7ac380b15b1dcee4cfa9bc Mon Sep 17 00:00:00 2001 From: Reiase Date: Thu, 20 Aug 2026 10:08:25 +0800 Subject: [PATCH 3/7] fix(pchronicle): gate ACTF parser by feature --- crates/persisting-pchronicle/src/formats/actf.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/persisting-pchronicle/src/formats/actf.rs b/crates/persisting-pchronicle/src/formats/actf.rs index 265e0813..42c7bb1c 100644 --- a/crates/persisting-pchronicle/src/formats/actf.rs +++ b/crates/persisting-pchronicle/src/formats/actf.rs @@ -106,6 +106,7 @@ pub struct ActfObservation { } impl ActfDocument { + #[cfg(any(test, feature = "lance-store"))] pub fn from_json_str(input: &str) -> InputResult { let document: Self = serde_json::from_str(input).map_err(|error| InputIssue::invalid(error.to_string()))?; From cc73d0dc893cf1f40af969f377d2e7a9bf2d5d3c Mon Sep 17 00:00:00 2001 From: Reiase Date: Thu, 20 Aug 2026 10:22:29 +0800 Subject: [PATCH 4/7] test(gateway): reject legacy AgenticMD fixture --- crates/persisting-gateway/tests/agenticmd_bridge.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/crates/persisting-gateway/tests/agenticmd_bridge.rs b/crates/persisting-gateway/tests/agenticmd_bridge.rs index b84565e8..5e1337d3 100644 --- a/crates/persisting-gateway/tests/agenticmd_bridge.rs +++ b/crates/persisting-gateway/tests/agenticmd_bridge.rs @@ -65,12 +65,12 @@ fn capture_turns_roundtrip_through_public_storyline_api() { } #[test] -fn legacy_fixture_still_imports_as_storyline() { - let story = decode_agenticmd(fixture()).expect("legacy AgenticMD import parse"); - assert_eq!(story.turns.len(), 2); - assert_eq!(story.turns[0].source, "user"); - assert_eq!(story.turns[0].message, json!("你好")); - assert_eq!(story.turns[1].source, "agent"); +fn legacy_fixture_without_authoritative_storyline_is_rejected() { + let error = decode_agenticmd(fixture()).unwrap_err(); + assert_eq!( + error.to_string(), + "missing authoritative Storyline metadata" + ); } #[test] From de5a0a2ab29510f7ef6ff8f0479375b3da5ff92c Mon Sep 17 00:00:00 2001 From: Reiase Date: Thu, 20 Aug 2026 10:39:34 +0800 Subject: [PATCH 5/7] test(gateway): reject legacy AgenticMD golden --- crates/persisting-gateway/tests/agenticmd_golden.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/crates/persisting-gateway/tests/agenticmd_golden.rs b/crates/persisting-gateway/tests/agenticmd_golden.rs index 66929b9f..17c34ce0 100644 --- a/crates/persisting-gateway/tests/agenticmd_golden.rs +++ b/crates/persisting-gateway/tests/agenticmd_golden.rs @@ -53,12 +53,12 @@ fn generated_agenticmd_preserves_golden_storyline_semantics() { } #[test] -fn checked_in_legacy_golden_remains_readable() { +fn checked_in_legacy_golden_without_authoritative_storyline_is_rejected() { let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) .join("tests/fixtures/agenticmd/demo-run-001.md"); - let story = decode_agenticmd(&std::fs::read_to_string(&fixture).unwrap()).unwrap(); - assert_eq!(story.session_id, "demo-run-001"); - assert_eq!(story.turns.len(), 2); - assert_eq!(story.turns[0].message, json!("你好")); - assert_eq!(story.turns[1].message, json!("你好！有什么可以帮你的？")); + let error = decode_agenticmd(&std::fs::read_to_string(&fixture).unwrap()).unwrap_err(); + assert_eq!( + error.to_string(), + "missing authoritative Storyline metadata" + ); } From a199720fea5c649a8dbbc54821c8f514423ef730 Mon Sep 17 00:00:00 2001 From: Reiase Date: Thu, 20 Aug 2026 11:20:24 +0800 Subject: [PATCH 6/7] fix(pchronicle): omit empty unknown metadata --- crates/persisting-pchronicle-cli/src/tests.rs | 6 ++-- .../src/formats/openai_corpus.rs | 35 +++++++++++++------ .../src/formats/unknown_fields.rs | 17 ++++++++- 3 files changed, 44 insertions(+), 14 deletions(-) diff --git a/crates/persisting-pchronicle-cli/src/tests.rs b/crates/persisting-pchronicle-cli/src/tests.rs index 8ac4e10c..0a228b8b 100644 --- a/crates/persisting-pchronicle-cli/src/tests.rs +++ b/crates/persisting-pchronicle-cli/src/tests.rs @@ -1414,8 +1414,10 @@ async fn export_filters_complete_trajectories_and_streams_finite_json() -> Resul let mut stderr = Vec::new(); run(cli, false, &mut stdout, &mut stderr).await?; - let rows: Value = serde_json::from_slice(&stdout)?; - let rows = rows.as_array().context("OpenAI export must be an array")?; + let document: Value = serde_json::from_slice(&stdout)?; + let rows = document["session_steps"] + .as_array() + .context("OpenAI export must contain a session_steps array")?; assert_eq!(rows.len(), 1); assert_eq!(rows[0]["session_id"], "training-002"); assert!(String::from_utf8(stderr)?.contains("trajectories=1")); diff --git a/crates/persisting-pchronicle/src/formats/openai_corpus.rs b/crates/persisting-pchronicle/src/formats/openai_corpus.rs index a9501bb8..36915ddb 100644 --- a/crates/persisting-pchronicle/src/formats/openai_corpus.rs +++ b/crates/persisting-pchronicle/src/formats/openai_corpus.rs @@ -19,8 +19,7 @@ use crate::formats::storyline::{ use crate::formats::unknown_fields::{ attach_carried_unknown_fields, normalize_openai_pointer, restore_json_pointer, take_unknown_fields_envelope, validate_unknown_fields_with, - write_foreign_unknown_fields_envelope, CarrierBinding, PointerWrite, SourceUnknownFields, - UnknownFieldLimits, + write_foreign_unknown_fields_envelope, CarrierBinding, PointerWrite, UnknownFieldLimits, }; use crate::{InputIssue, InputResult, Result}; @@ -130,14 +129,6 @@ fn capture_openai_unknowns( root_unknown: &Map, records: &[(usize, Value)], ) -> InputResult<()> { - story - .unknown_fields - .sources - .entry("openai-msg".into()) - .or_insert_with(|| SourceUnknownFields { - source_document_id: source_document_id.to_string(), - fields: Default::default(), - }); insert_openai_map(story, source_document_id, "", root_unknown)?; for (ordinal, record) in records { let row = record.as_object().ok_or_else(|| { @@ -338,7 +329,7 @@ pub(crate) fn has_openai_provenance(story: &StorylineDocument) -> bool { .is_some() } -/// Explicitly synthesize an OpenAI message row array from Storyline semantics. +/// Explicitly synthesize an OpenAI `session_steps` envelope from Storyline semantics. /// /// This is a cross-format projection, not a lossless recovery operation. Use /// [`recover_openai_msg_files`] when the Storylines originated from an OpenAI @@ -1155,6 +1146,28 @@ mod tests { stories[0].unknown_key_counts["openai-msg"]["/session_steps/*/messages/*/0"], 1 ); + assert_eq!( + crate::formats::unknown_fields::compute_unknown_key_counts(&stories[0].unknown_fields) + .unwrap(), + stories[0].unknown_key_counts + ); + stories[0].validate().unwrap(); + } + + #[test] + fn openai_known_only_rows_do_not_create_empty_unknown_metadata() { + let input = json!([{ + "session_id": "s", + "step_id": 1, + "messages": [{"role": "user", "content": "hi"}], + "response": {"role": "assistant", "content": "ok"} + }]); + + let stories = parse_openai_msg_corpus_value(&input, "corpus.json").unwrap(); + + assert!(stories[0].unknown_fields.sources.is_empty()); + assert!(stories[0].unknown_key_counts.is_empty()); + stories[0].validate().unwrap(); } #[cfg(feature = "lance-store")] use crate::store::StorylineLanceStore; diff --git a/crates/persisting-pchronicle/src/formats/unknown_fields.rs b/crates/persisting-pchronicle/src/formats/unknown_fields.rs index b0e63e94..8b092f53 100644 --- a/crates/persisting-pchronicle/src/formats/unknown_fields.rs +++ b/crates/persisting-pchronicle/src/formats/unknown_fields.rs @@ -452,10 +452,10 @@ where { let mut counts = UnknownKeyCounts::new(); for (source, source_fields) in &fields.sources { - let source_counts = counts.entry(source.clone()).or_default(); for pointer in source_fields.fields.keys() { validate_json_pointer(pointer)?; let normalized_pointer = normalize(source, pointer)?; + let source_counts = counts.entry(source.clone()).or_default(); let count = source_counts.entry(normalized_pointer).or_default(); *count = count.saturating_add(1); } @@ -755,6 +755,21 @@ mod tests { use crate::formats::StorylineDocument; use serde_json::json; + #[test] + fn empty_source_does_not_create_a_key_count_entry() { + let fields = StorylineUnknownFields { + sources: BTreeMap::from([( + "openai-msg".into(), + SourceUnknownFields { + source_document_id: "source.json".into(), + fields: BTreeMap::new(), + }, + )]), + }; + + assert!(compute_unknown_key_counts(&fields).unwrap().is_empty()); + } + #[test] fn envelope_distributes_foreign_sources_by_carrier() { let mut raw = json!({ From b6c9b3eb6e059a6d8b0e07bc2d9fbf540e9f083a Mon Sep 17 00:00:00 2001 From: Reiase Date: Thu, 20 Aug 2026 12:42:37 +0800 Subject: [PATCH 7/7] test(pchronicle): assert canonical reimport stability --- .../tests/import_export_roundtrip.rs | 35 +++++++++++++++++-- 1 file changed, 32 insertions(+), 3 deletions(-) diff --git a/crates/persisting-pchronicle-cli/tests/import_export_roundtrip.rs b/crates/persisting-pchronicle-cli/tests/import_export_roundtrip.rs index 4241998f..d56faf8a 100644 --- a/crates/persisting-pchronicle-cli/tests/import_export_roundtrip.rs +++ b/crates/persisting-pchronicle-cli/tests/import_export_roundtrip.rs @@ -80,7 +80,7 @@ async fn import_export_roundtrip_is_byte_identical_and_reimportable() -> Result< } #[tokio::test] -async fn forced_storyline_roundtrip_is_canonical_json_byte_identical() -> Result<()> { +async fn forced_storyline_roundtrip_is_canonical_and_reimport_stable() -> Result<()> { let temp = tempfile::tempdir()?; for fixture in EXAMPLE_FIXTURES { let format = fixture.name; @@ -115,10 +115,39 @@ async fn forced_storyline_roundtrip_is_canonical_json_byte_identical() -> Result .await?; assert!(exported_output.stderr_text()?.contains("exact=false")); + let reimported = temp.path().join(format!("{format}-storyline-reimported")); + run_cli([ + "import", + "--from", + exported.to_str().unwrap(), + "--output", + reimported.to_str().unwrap(), + "--format", + format, + ]) + .await?; + + let reexported = temp + .path() + .join(format!("{format}-storyline-reexport.json")); + let reexported_output = run_cli([ + "export", + "--from", + reimported.to_str().unwrap(), + "--output", + reexported.to_str().unwrap(), + "--format", + format, + "--where", + "TRUE", + ]) + .await?; + assert!(reexported_output.stderr_text()?.contains("exact=false")); + assert_eq!( + canonical_json_bytes(&reexported)?, canonical_json_bytes(&exported)?, - canonical_json_bytes(&input)?, - "Storyline round-trip canonical JSON differs for {format}" + "Storyline canonical JSON is not reimport-stable for {format}" ); } Ok(())