From c79e102b94e75ccd8158f929ef88099412f0aac2 Mon Sep 17 00:00:00 2001 From: Vishwanath Seshagiri Date: Tue, 16 Jun 2026 21:55:13 -0700 Subject: [PATCH 1/7] Add agent-trace-core: Rust core via PyO3 + maturin Rust crate with trace models, NDJSON + SHA-256 hash chain, and Claude Code JSONL import. Builds as both a pure-Rust library and an abi3 Python extension (agent_trace_core) via maturin/PyO3. Output matches the existing agent_trace.jsonl_import format, so the Python tooling reads it unchanged. --- crates/agent-trace-core/.gitignore | 2 + crates/agent-trace-core/Cargo.toml | 27 ++ crates/agent-trace-core/README.md | 76 +++++ crates/agent-trace-core/pyproject.toml | 21 ++ crates/agent-trace-core/src/import.rs | 445 +++++++++++++++++++++++++ crates/agent-trace-core/src/lib.rs | 74 ++++ crates/agent-trace-core/src/models.rs | 298 +++++++++++++++++ 7 files changed, 943 insertions(+) create mode 100644 crates/agent-trace-core/.gitignore create mode 100644 crates/agent-trace-core/Cargo.toml create mode 100644 crates/agent-trace-core/README.md create mode 100644 crates/agent-trace-core/pyproject.toml create mode 100644 crates/agent-trace-core/src/import.rs create mode 100644 crates/agent-trace-core/src/lib.rs create mode 100644 crates/agent-trace-core/src/models.rs diff --git a/crates/agent-trace-core/.gitignore b/crates/agent-trace-core/.gitignore new file mode 100644 index 0000000..96ef6c0 --- /dev/null +++ b/crates/agent-trace-core/.gitignore @@ -0,0 +1,2 @@ +/target +Cargo.lock diff --git a/crates/agent-trace-core/Cargo.toml b/crates/agent-trace-core/Cargo.toml new file mode 100644 index 0000000..4576190 --- /dev/null +++ b/crates/agent-trace-core/Cargo.toml @@ -0,0 +1,27 @@ +[package] +name = "agent-trace-core" +version = "0.1.0" +edition = "2021" +description = "Rust core for agent-trace: trace models, NDJSON with a SHA-256 hash chain, and Claude Code session import." +license = "MIT" +repository = "https://github.com/Siddhant-K-code/agent-trace" + +# cdylib -> the Python extension module (built by maturin) +# rlib -> usable as a plain Rust library (the "SDK" surface) +[lib] +name = "agent_trace_core" +crate-type = ["cdylib", "rlib"] + +[dependencies] +serde = { version = "1", features = ["derive"] } +# preserve_order keeps JSON object keys in insertion order so our output +# matches the field order the Python dataclasses emit. +serde_json = { version = "1", features = ["preserve_order"] } +sha2 = "0.10" +chrono = { version = "0.4", default-features = false, features = ["std"] } +pyo3 = { version = "0.24", features = ["abi3-py310"] } + +[features] +# maturin enables this when building the wheel. Leaving it off lets `cargo test` +# link libpython normally so the pure-Rust unit tests run without a wheel. +extension-module = ["pyo3/extension-module"] diff --git a/crates/agent-trace-core/README.md b/crates/agent-trace-core/README.md new file mode 100644 index 0000000..e932664 --- /dev/null +++ b/crates/agent-trace-core/README.md @@ -0,0 +1,76 @@ +# agent-trace-core + +The Rust core for [agent-trace](../../README.md). One crate, two surfaces: + +- **Rust library** (`rlib`) — `models` (trace types, NDJSON, SHA-256 hash chain) + and `import` (Claude Code JSONL import). Use it from any Rust program. +- **Python extension** (`cdylib`, built with [maturin](https://www.maturin.rs/) + + [PyO3](https://pyo3.rs/)) — the module `agent_trace_core`, exposing the same + functionality to Python with no per-event Python object allocation. Parsing + and verifying large traces stays flat in memory. + +## Why Rust + +The hot paths — parsing NDJSON, walking the hash chain, importing multi-MB +Claude session logs — allocate one short-lived object per event in Python. In +Rust those are stack values dropped deterministically, so memory stays bounded +regardless of trace size, and there's no GC to chase. + +## Build + +```bash +# from this directory +maturin develop --release # build + install into the active venv +maturin build --release # produce a wheel under ../../target/wheels/ +``` + +## Python API + +```python +import agent_trace_core as core + +# Discover Claude Code sessions under ~/.claude/projects/ +for s in core.discover_claude_sessions(): # claude_dir defaults to ~/.claude + print(s["session_id"], s["size_kb"], s["project"]) + +# Import a session into .agent-traces// (meta.json + events.ndjson) +summary = core.import_claude_jsonl(path, trace_dir=".agent-traces") +print(summary) # {'session_id': ..., 'tool_calls': 6, 'llm_requests': 14, ...} + +# Parse / verify NDJSON +lines = core.parse_ndjson(open("events.ndjson").read()) # canonicalized lines +ok = core.verify_hash_chain(open("events.ndjson").read()) # tamper check +``` + +## Rust API + +```rust +use agent_trace_core::import::{import_jsonl, discover_claude_sessions}; +use agent_trace_core::models::{parse_ndjson, verify_hash_chain, TraceEvent}; + +let summary = import_jsonl("session.jsonl", ".agent-traces")?; +let intact = verify_hash_chain(&std::fs::read_to_string("events.ndjson")?); +``` + +## Compatibility with the Python implementation + +The serialized output matches `agent_trace.models` field-for-field (compact +event lines, conditional key dropping, pretty meta), so files written here are +read by the existing Python tooling and round-trip cleanly. + +One intentional difference: imported events get sequential `event_id`s +(`000000000000`, `000000000001`, …) instead of Python's random UUID fragments. +`event_id` isn't used to correlate imported events, and sequential ids make the +hash chain reproducible for tests. The chain itself is always internally +consistent (`verify_hash_chain` passes). + +## Tests + +The unit tests are pure Rust, but the test binary links PyO3, so the loader +needs `libpython` on its path: + +```bash +LIBDIR=$(python3 -c "import sysconfig; print(sysconfig.get_config_var('LIBDIR'))") +DYLD_FALLBACK_LIBRARY_PATH="$LIBDIR" cargo test # macOS +LD_LIBRARY_PATH="$LIBDIR" cargo test # Linux +``` diff --git a/crates/agent-trace-core/pyproject.toml b/crates/agent-trace-core/pyproject.toml new file mode 100644 index 0000000..0eedc34 --- /dev/null +++ b/crates/agent-trace-core/pyproject.toml @@ -0,0 +1,21 @@ +[build-system] +requires = ["maturin>=1.7,<2.0"] +build-backend = "maturin" + +[project] +name = "agent-trace-core" +description = "Rust core for agent-trace: trace models, NDJSON hash chain, and Claude Code import." +readme = "README.md" +license = "MIT" +requires-python = ">=3.10" +classifiers = [ + "Programming Language :: Rust", + "Programming Language :: Python :: Implementation :: CPython", +] +dynamic = ["version"] + +[tool.maturin] +# Pure-Rust extension; the module name must match the #[pymodule] in lib.rs. +module-name = "agent_trace_core" +# abi3 wheel -> one build works across CPython 3.10+. +features = ["pyo3/abi3-py310", "extension-module"] diff --git a/crates/agent-trace-core/src/import.rs b/crates/agent-trace-core/src/import.rs new file mode 100644 index 0000000..bccc792 --- /dev/null +++ b/crates/agent-trace-core/src/import.rs @@ -0,0 +1,445 @@ +//! Import Claude Code native JSONL session logs. + +use std::fs; +use std::io; +use std::path::{Path, PathBuf}; + +use chrono::DateTime; +use serde_json::{json, Map, Value}; + +use crate::models::{ + compact_number, sha256_hex, SessionMeta, TraceEvent, ASSISTANT_RESPONSE, SESSION_END, + TOOL_CALL, TOOL_RESULT, USER_PROMPT, +}; + +/// Summary returned after importing a session. +pub struct ImportSummary { + pub session_id: String, + pub tool_calls: u64, + pub llm_requests: u64, + pub total_tokens: u64, + pub events: usize, +} + +/// Info about a discovered Claude Code session log. +pub struct SessionInfo { + pub path: PathBuf, + pub project: String, + pub session_id: String, + pub size_kb: u64, +} + +/// Expand a leading `~` to `$HOME`. +pub fn expanduser(p: &str) -> PathBuf { + if let Some(rest) = p.strip_prefix("~") { + if let Ok(home) = std::env::var("HOME") { + let rest = rest.strip_prefix('/').unwrap_or(rest); + return Path::new(&home).join(rest); + } + } + PathBuf::from(p) +} + +/// Convert an ISO 8601 timestamp to Unix epoch seconds (0.0 on failure). +fn parse_iso_timestamp(ts: &str) -> f64 { + if ts.is_empty() { + return 0.0; + } + match DateTime::parse_from_rfc3339(ts) { + Ok(dt) => dt.timestamp() as f64 + (dt.timestamp_subsec_nanos() as f64) / 1e9, + Err(_) => 0.0, + } +} + +/// Truncate to at most `n` Unicode chars. +fn take_chars(s: &str, n: usize) -> String { + s.chars().take(n).collect() +} + +fn char_len(s: &str) -> usize { + s.chars().count() +} + +/// Extract text from message content (a string, or a list of content blocks). +fn extract_text(content: &Value) -> String { + match content { + Value::String(s) => s.clone(), + Value::Array(blocks) => { + let mut parts = Vec::new(); + for block in blocks { + if block.get("type").and_then(Value::as_str) == Some("text") { + parts.push(block.get("text").and_then(Value::as_str).unwrap_or("").to_string()); + } + } + parts.join("\n") + } + _ => String::new(), + } +} + +struct ToolCall { + id: String, + name: String, + input: Value, + caller: Value, +} + +fn extract_tool_calls(content: &[Value]) -> Vec { + let mut calls = Vec::new(); + for block in content { + if !block.is_object() { + continue; + } + if block.get("type").and_then(Value::as_str) == Some("tool_use") { + calls.push(ToolCall { + id: block.get("id").and_then(Value::as_str).unwrap_or("").to_string(), + name: block.get("name").and_then(Value::as_str).unwrap_or("").to_string(), + input: block.get("input").cloned().unwrap_or_else(|| json!({})), + caller: block.get("caller").cloned().unwrap_or_else(|| json!({})), + }); + } + } + calls +} + +struct ToolResult { + tool_use_id: String, + content: String, +} + +fn extract_tool_results(content: &[Value]) -> Vec { + let mut results = Vec::new(); + for block in content { + if !block.is_object() { + continue; + } + if block.get("type").and_then(Value::as_str) == Some("tool_result") { + let mut text_parts = Vec::new(); + match block.get("content") { + Some(Value::String(s)) => text_parts.push(s.clone()), + Some(Value::Array(subs)) => { + for sub in subs { + if sub.get("type").and_then(Value::as_str) == Some("text") { + text_parts.push(sub.get("text").and_then(Value::as_str).unwrap_or("").to_string()); + } + } + } + _ => {} + } + results.push(ToolResult { + tool_use_id: block.get("tool_use_id").and_then(Value::as_str).unwrap_or("").to_string(), + content: text_parts.join("\n"), + }); + } + } + results +} + +fn obj_get_str<'a>(v: &'a Value, key: &str) -> &'a str { + v.get(key).and_then(Value::as_str).unwrap_or("") +} + +fn usage_u64(usage: &Value, key: &str) -> u64 { + usage.get(key).and_then(Value::as_u64).unwrap_or(0) +} + +/// Import a Claude Code JSONL session log into `//`. +pub fn import_jsonl(path: &str, trace_dir: &str) -> io::Result { + let path = expanduser(path); + let path = fs::canonicalize(&path).unwrap_or(path); + if !path.exists() { + return Err(io::Error::new( + io::ErrorKind::NotFound, + format!("Session log not found: {}", path.display()), + )); + } + + let text = fs::read_to_string(&path)?; + + // First pass: collect entries + session-level metadata. + let mut session_id = String::new(); + let mut git_branch = String::new(); + let mut version = String::new(); + let mut first_ts = 0.0_f64; + let mut last_ts = 0.0_f64; + let mut entries: Vec = Vec::new(); + + for line in text.lines() { + let line = line.trim(); + if line.is_empty() { + continue; + } + let raw: Value = match serde_json::from_str(line) { + Ok(v) => v, + Err(_) => continue, + }; + if raw.get("type").and_then(Value::as_str) == Some("queue-operation") { + continue; + } + + let ts = parse_iso_timestamp(obj_get_str(&raw, "timestamp")); + if first_ts == 0.0 && ts > 0.0 { + first_ts = ts; + } + if ts > 0.0 { + last_ts = ts; + } + if session_id.is_empty() { + session_id = obj_get_str(&raw, "sessionId").to_string(); + git_branch = obj_get_str(&raw, "gitBranch").to_string(); + version = obj_get_str(&raw, "version").to_string(); + } + entries.push(raw); + } + + if session_id.is_empty() { + session_id = path.file_stem().and_then(|s| s.to_str()).unwrap_or("session").to_string(); + } + + let mut meta = SessionMeta::new( + session_id.clone(), + first_ts, + "claude-code", + format!( + "imported from {} (branch: {}, v{})", + path.file_name().and_then(|s| s.to_str()).unwrap_or(""), + git_branch, + version + ), + ); + + // Second pass: convert entries to events. event_id uses a per-import + // counter (unique within the session); see README for why this differs + // from the Python random ids. + let mut events: Vec = Vec::new(); + let mut counter: u64 = 0; + let new_id = |c: &mut u64| -> String { + let id = format!("{:012x}", *c); + *c += 1; + id + }; + + for raw in &entries { + let entry_type = obj_get_str(raw, "type"); + let ts = parse_iso_timestamp(obj_get_str(raw, "timestamp")); + // Match Python's `msg = raw.get("message", {})`: a missing message is an + // empty object (entry still processed, e.g. `system`/turn_duration); a + // present-but-non-object message skips the entry. + let empty = json!({}); + let msg = match raw.get("message") { + Some(m) if m.is_object() => m, + Some(_) => continue, + None => &empty, + }; + let content = msg.get("content").cloned().unwrap_or(Value::Null); + let usage = msg.get("usage").cloned().unwrap_or(Value::Null); + let is_sidechain = raw.get("isSidechain").and_then(Value::as_bool).unwrap_or(false); + + if entry_type == "user" { + let text_val = extract_text(&content); + + if let Value::Array(blocks) = &content { + let tool_results = extract_tool_results(blocks); + for tr in &tool_results { + let preview = if tr.content.is_empty() { + String::new() + } else if char_len(&tr.content) > 2000 { + format!("{}...", take_chars(&tr.content, 2000)) + } else { + tr.content.clone() + }; + events.push(TraceEvent::new( + TOOL_RESULT, + ts, + new_id(&mut counter), + session_id.clone(), + json!({ "tool_use_id": tr.tool_use_id, "content_preview": preview }), + )); + } + + if tool_results.is_empty() { + if let Some(tr_data) = raw.get("toolUseResult") { + if tr_data.is_object() { + let stdout = obj_get_str(tr_data, "stdout"); + let stderr = obj_get_str(tr_data, "stderr"); + if !stdout.is_empty() || !stderr.is_empty() { + let mut result_text = take_chars(stdout, 500); + if !stderr.is_empty() { + result_text.push_str(&format!(" [stderr: {}]", take_chars(stderr, 200))); + } + events.push(TraceEvent::new( + TOOL_RESULT, + ts, + new_id(&mut counter), + session_id.clone(), + json!({ "result": result_text, "content_types": ["text"] }), + )); + } + } + } + } + } + + if !text_val.is_empty() && !text_val.starts_with('{') { + events.push(TraceEvent::new( + USER_PROMPT, + ts, + new_id(&mut counter), + session_id.clone(), + json!({ "prompt": take_chars(&text_val, 2000) }), + )); + } + } else if entry_type == "assistant" { + let text_val = extract_text(&content); + + if let Value::Array(blocks) = &content { + for tc in extract_tool_calls(blocks) { + let mut data = Map::new(); + data.insert("tool_name".into(), json!(tc.name)); + data.insert("arguments".into(), tc.input.clone()); + data.insert("request_id".into(), json!(tc.id)); + if is_sidechain { + data.insert("is_sidechain".into(), json!(true)); + } + let caller_type = obj_get_str(&tc.caller, "type"); + if !caller_type.is_empty() { + data.insert("caller_type".into(), json!(caller_type)); + } + if tc.name == "Agent" { + let subagent = obj_get_str(&tc.input, "subagent_type"); + if !subagent.is_empty() { + data.insert("subagent_type".into(), json!(subagent)); + } + } + events.push(TraceEvent::new( + TOOL_CALL, + ts, + new_id(&mut counter), + session_id.clone(), + Value::Object(data), + )); + meta.tool_calls += 1; + } + } + + if !text_val.is_empty() { + events.push(TraceEvent::new( + ASSISTANT_RESPONSE, + ts, + new_id(&mut counter), + session_id.clone(), + json!({ "text": take_chars(&text_val, 2000), "model": model }), + )); + } + + if usage.is_object() { + meta.total_tokens += usage_u64(&usage, "input_tokens") + + usage_u64(&usage, "output_tokens") + + usage_u64(&usage, "cache_creation_input_tokens") + + usage_u64(&usage, "cache_read_input_tokens"); + meta.llm_requests += 1; + } + } else if entry_type == "system" { + if obj_get_str(raw, "subtype") == "turn_duration" { + let duration_ms = raw.get("durationMs").and_then(Value::as_f64).unwrap_or(0.0); + if duration_ms != 0.0 { + meta.total_duration_ms += duration_ms; + } + } + } + } + + // Finalize. + meta.ended_at = Some(if last_ts > 0.0 { last_ts } else { meta.started_at }); + let ended = meta.ended_at.unwrap(); + if meta.total_duration_ms == 0.0 && ended > meta.started_at { + meta.total_duration_ms = (ended - meta.started_at) * 1000.0; + } + + events.push(TraceEvent::new( + SESSION_END, + ended, + new_id(&mut counter), + session_id.clone(), + json!({ + "duration_ms": compact_number(meta.total_duration_ms), + "tool_calls": meta.tool_calls, + "llm_requests": meta.llm_requests, + "total_tokens": meta.total_tokens, + "source": path.to_string_lossy(), + }), + )); + + // Write the session directory. + let session_dir = Path::new(trace_dir).join(&session_id); + fs::create_dir_all(&session_dir)?; + fs::write(session_dir.join("meta.json"), meta.to_json())?; + + // Build the events file with the SHA-256 hash chain. + let mut ndjson = String::new(); + let mut prev_line = String::new(); + for ev in events.iter_mut() { + if !prev_line.is_empty() { + ev.prev_hash = sha256_hex(&prev_line); + } + let line = ev.to_json(); + ndjson.push_str(&line); + ndjson.push('\n'); + prev_line = line; + } + fs::write(session_dir.join("events.ndjson"), ndjson)?; + + Ok(ImportSummary { + session_id: meta.session_id, + tool_calls: meta.tool_calls, + llm_requests: meta.llm_requests, + total_tokens: meta.total_tokens, + events: events.len(), + }) +} + +/// Decode Claude Code's encoded project directory name (matches Python). +fn decode_project_path(encoded: &str) -> String { + if !encoded.starts_with('-') { + return encoded.to_string(); + } + encoded.replace('-', "/") +} + +/// Discover all Claude Code session JSONL files under `/projects/`. +pub fn discover_claude_sessions(claude_dir: &str) -> Vec { + let claude_dir = expanduser(claude_dir); + let projects_dir = claude_dir.join("projects"); + let mut sessions = Vec::new(); + let mut project_dirs: Vec = match fs::read_dir(&projects_dir) { + Ok(rd) => rd.filter_map(|e| e.ok().map(|e| e.path())).filter(|p| p.is_dir()).collect(), + Err(_) => return sessions, + }; + project_dirs.sort(); + + for project_dir in project_dirs { + let name = project_dir.file_name().and_then(|s| s.to_str()).unwrap_or(""); + let project_name = decode_project_path(name); + + let mut jsonl_files: Vec = match fs::read_dir(&project_dir) { + Ok(rd) => rd + .filter_map(|e| e.ok().map(|e| e.path())) + .filter(|p| p.extension().and_then(|x| x.to_str()) == Some("jsonl")) + .collect(), + Err(_) => continue, + }; + jsonl_files.sort(); + + for f in jsonl_files { + let size_kb = f.metadata().map(|m| m.len() / 1024).unwrap_or(0); + let session_id = f.file_stem().and_then(|s| s.to_str()).unwrap_or("").to_string(); + sessions.push(SessionInfo { + path: f, + project: project_name.clone(), + session_id, + size_kb, + }); + } + } + sessions +} diff --git a/crates/agent-trace-core/src/lib.rs b/crates/agent-trace-core/src/lib.rs new file mode 100644 index 0000000..ccd057c --- /dev/null +++ b/crates/agent-trace-core/src/lib.rs @@ -0,0 +1,74 @@ +//! `agent_trace_core` — the Rust core for agent-trace. +//! +//! Two surfaces from one crate: +//! * a plain Rust library (`rlib`) — see [`models`] and [`import`]; +//! * a Python extension module (`cdylib`, built by maturin) exposing the +//! same functionality with no per-event Python allocations, so importing +//! and verifying large traces stays flat in memory. + +pub mod import; +pub mod models; + +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyList}; + +/// Import a Claude Code JSONL session log into `//`. +/// Returns a summary dict: session_id, tool_calls, llm_requests, total_tokens, events. +#[pyfunction] +#[pyo3(signature = (path, trace_dir=None))] +fn import_claude_jsonl(py: Python<'_>, path: String, trace_dir: Option) -> PyResult { + let trace_dir = trace_dir.unwrap_or_else(|| ".agent-traces".to_string()); + let summary = import::import_jsonl(&path, &trace_dir) + .map_err(|e| pyo3::exceptions::PyOSError::new_err(e.to_string()))?; + let d = PyDict::new(py); + d.set_item("session_id", summary.session_id)?; + d.set_item("tool_calls", summary.tool_calls)?; + d.set_item("llm_requests", summary.llm_requests)?; + d.set_item("total_tokens", summary.total_tokens)?; + d.set_item("events", summary.events)?; + Ok(d.into()) +} + +/// Discover Claude Code session logs under `/projects/`. +/// Returns a list of dicts: path, project, session_id, size_kb. +#[pyfunction] +#[pyo3(signature = (claude_dir=None))] +fn discover_claude_sessions(py: Python<'_>, claude_dir: Option) -> PyResult { + let claude_dir = claude_dir.unwrap_or_else(|| "~/.claude".to_string()); + let sessions = import::discover_claude_sessions(&claude_dir); + let list = PyList::empty(py); + for s in sessions { + let d = PyDict::new(py); + d.set_item("path", s.path.to_string_lossy().into_owned())?; + d.set_item("project", s.project)?; + d.set_item("session_id", s.session_id)?; + d.set_item("size_kb", s.size_kb)?; + list.append(d)?; + } + Ok(list.into()) +} + +/// Parse an NDJSON trace, returning canonicalized event lines (round-tripped +/// through the Rust model). Useful for validation / re-serialization. +#[pyfunction] +fn parse_ndjson(text: &str) -> PyResult> { + let events = models::parse_ndjson(text) + .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?; + Ok(events.iter().map(|e| e.to_json()).collect()) +} + +/// Return true if the NDJSON trace's SHA-256 hash chain is intact. +#[pyfunction] +fn verify_hash_chain(text: &str) -> bool { + models::verify_hash_chain(text) +} + +#[pymodule] +fn agent_trace_core(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add("__version__", env!("CARGO_PKG_VERSION"))?; + m.add_function(wrap_pyfunction!(import_claude_jsonl, m)?)?; + m.add_function(wrap_pyfunction!(discover_claude_sessions, m)?)?; + m.add_function(wrap_pyfunction!(parse_ndjson, m)?)?; + m.add_function(wrap_pyfunction!(verify_hash_chain, m)?)?; + Ok(()) +} diff --git a/crates/agent-trace-core/src/models.rs b/crates/agent-trace-core/src/models.rs new file mode 100644 index 0000000..a8342e1 --- /dev/null +++ b/crates/agent-trace-core/src/models.rs @@ -0,0 +1,298 @@ +//! Trace data model + NDJSON serialization. + +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use sha2::{Digest, Sha256}; + +pub const SESSION_END: &str = "session_end"; +pub const TOOL_CALL: &str = "tool_call"; +pub const TOOL_RESULT: &str = "tool_result"; +pub const USER_PROMPT: &str = "user_prompt"; +pub const ASSISTANT_RESPONSE: &str = "assistant_response"; + +fn is_false(b: &bool) -> bool { + !*b +} + +fn empty_object() -> Value { + Value::Object(serde_json::Map::new()) +} + +/// A single event in a trace. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TraceEvent { + pub event_type: String, + pub timestamp: f64, + pub event_id: String, + #[serde(default, skip_serializing_if = "String::is_empty")] + pub session_id: String, + #[serde(default, skip_serializing_if = "String::is_empty")] + pub parent_id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub duration_ms: Option, + #[serde(default = "empty_object")] + pub data: Value, + #[serde(default, skip_serializing_if = "String::is_empty")] + pub prev_hash: String, + #[serde(default, skip_serializing_if = "is_false")] + pub redacted: bool, +} + +impl TraceEvent { + pub fn new(event_type: &str, timestamp: f64, event_id: String, session_id: String, data: Value) -> Self { + TraceEvent { + event_type: event_type.to_string(), + timestamp, + event_id, + session_id, + parent_id: String::new(), + duration_ms: None, + data, + prev_hash: String::new(), + redacted: false, + } + } + + /// Compact JSON, byte-for-byte equivalent to the Python `to_json` output + /// (`separators=(",", ":")`, conditional key dropping, field ordering). + pub fn to_json(&self) -> String { + serde_json::to_string(self).expect("TraceEvent serializes") + } + + pub fn from_json(line: &str) -> serde_json::Result { + serde_json::from_str(line) + } +} + +/// Session metadata. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SessionMeta { + pub session_id: String, + pub started_at: f64, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub ended_at: Option, + #[serde(default, skip_serializing_if = "String::is_empty")] + pub agent_name: String, + #[serde(default, skip_serializing_if = "String::is_empty")] + pub command: String, + #[serde(default, skip_serializing_if = "is_zero_u64")] + pub tool_calls: u64, + #[serde(default, skip_serializing_if = "is_zero_u64")] + pub llm_requests: u64, + #[serde(default, skip_serializing_if = "is_zero_u64")] + pub errors: u64, + #[serde(default, skip_serializing_if = "is_zero_u64")] + pub total_tokens: u64, + #[serde(default, skip_serializing_if = "is_zero_f64", serialize_with = "serialize_compact_f64")] + pub total_duration_ms: f64, + #[serde(default, skip_serializing_if = "String::is_empty")] + pub parent_session_id: String, + #[serde(default, skip_serializing_if = "String::is_empty")] + pub parent_event_id: String, + #[serde(default, skip_serializing_if = "is_zero_i64")] + pub depth: i64, + #[serde(default, skip_serializing_if = "String::is_empty")] + pub team: String, + #[serde(default, skip_serializing_if = "String::is_empty")] + pub workspace_id: String, + #[serde(default = "empty_object")] + pub attribution: Value, + #[serde(default, skip_serializing_if = "is_false")] + pub redacted: bool, +} + +/// Largest integer exactly representable as f64. +const MAX_EXACT_F64_INT: f64 = 9_007_199_254_740_992.0; + +/// JSON number rendering a whole f64 as an integer (`50058`, not `50058.0`). +pub fn compact_number(v: f64) -> Value { + if v.is_finite() && v.fract() == 0.0 && v.abs() < MAX_EXACT_F64_INT { + Value::from(v as i64) + } else { + Value::from(v) + } +} + +fn serialize_compact_f64(v: &f64, s: S) -> Result +where + S: serde::Serializer, +{ + if v.is_finite() && v.fract() == 0.0 && v.abs() < MAX_EXACT_F64_INT { + s.serialize_i64(*v as i64) + } else { + s.serialize_f64(*v) + } +} + +fn is_zero_u64(n: &u64) -> bool { + *n == 0 +} +fn is_zero_i64(n: &i64) -> bool { + *n == 0 +} +fn is_zero_f64(n: &f64) -> bool { + *n == 0.0 +} + +impl SessionMeta { + pub fn new(session_id: String, started_at: f64, agent_name: &str, command: String) -> Self { + SessionMeta { + session_id, + started_at, + ended_at: None, + agent_name: agent_name.to_string(), + command, + tool_calls: 0, + llm_requests: 0, + errors: 0, + total_tokens: 0, + total_duration_ms: 0.0, + parent_session_id: String::new(), + parent_event_id: String::new(), + depth: 0, + team: String::new(), + workspace_id: String::new(), + attribution: empty_object(), + redacted: false, + } + } + + pub fn to_json(&self) -> String { + serde_json::to_string_pretty(self).expect("SessionMeta serializes") + } + + pub fn from_json(text: &str) -> serde_json::Result { + serde_json::from_str(text) + } +} + +/// Lowercase hex SHA-256 of `s`. +pub fn sha256_hex(s: &str) -> String { + let mut hasher = Sha256::new(); + hasher.update(s.as_bytes()); + let digest = hasher.finalize(); + let mut out = String::with_capacity(64); + for b in digest { + out.push_str(&format!("{:02x}", b)); + } + out +} + +/// Serialize a list of events into an NDJSON string, filling each event's +/// `prev_hash` with the SHA-256 of the previous line (empty for the first). +/// The events are mutated in place so the returned chain is self-consistent. +pub fn write_ndjson(events: &mut [TraceEvent]) -> String { + let mut out = String::new(); + let mut prev_line = String::new(); + for ev in events.iter_mut() { + if ev.prev_hash.is_empty() && !prev_line.is_empty() { + ev.prev_hash = sha256_hex(&prev_line); + } + let line = ev.to_json(); + out.push_str(&line); + out.push('\n'); + prev_line = line; + } + out +} + +/// Parse an NDJSON string into events, skipping blank lines. +pub fn parse_ndjson(text: &str) -> serde_json::Result> { + let mut events = Vec::new(); + for line in text.lines() { + let line = line.trim(); + if line.is_empty() { + continue; + } + events.push(TraceEvent::from_json(line)?); + } + Ok(events) +} + +/// Verify that every line's `prev_hash` equals the SHA-256 of the previous +/// line. Returns true for an intact chain (and for empty input). +pub fn verify_hash_chain(text: &str) -> bool { + let mut prev_line: Option<&str> = None; + for line in text.lines() { + let line = line.trim(); + if line.is_empty() { + continue; + } + let ev: TraceEvent = match TraceEvent::from_json(line) { + Ok(e) => e, + Err(_) => return false, + }; + let expected = match prev_line { + None => String::new(), + Some(p) => sha256_hex(p), + }; + if ev.prev_hash != expected { + return false; + } + prev_line = Some(line); + } + true +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn event_to_json_drops_empties_and_keeps_data() { + let ev = TraceEvent::new(USER_PROMPT, 1.5, "abc".into(), "sess".into(), json!({"prompt": "hi"})); + // parent_id, duration_ms, prev_hash, redacted all omitted; data kept. + assert_eq!( + ev.to_json(), + r#"{"event_type":"user_prompt","timestamp":1.5,"event_id":"abc","session_id":"sess","data":{"prompt":"hi"}}"# + ); + } + + #[test] + fn empty_data_is_still_emitted() { + let ev = TraceEvent::new(SESSION_END, 0.0, "id".into(), "".into(), empty_object()); + // session_id empty -> dropped; data:{} -> kept; timestamp 0.0 -> kept. + assert_eq!( + ev.to_json(), + r#"{"event_type":"session_end","timestamp":0.0,"event_id":"id","data":{}}"# + ); + } + + #[test] + fn hash_chain_round_trips_and_verifies() { + let mut events = vec![ + TraceEvent::new(USER_PROMPT, 1.0, "a".into(), "s".into(), json!({"prompt": "x"})), + TraceEvent::new(ASSISTANT_RESPONSE, 2.0, "b".into(), "s".into(), json!({"text": "y"})), + TraceEvent::new(TOOL_CALL, 3.0, "c".into(), "s".into(), json!({"tool_name": "Bash"})), + ]; + let ndjson = write_ndjson(&mut events); + assert!(verify_hash_chain(&ndjson)); + assert!(events[0].prev_hash.is_empty()); + assert_eq!(events[1].prev_hash.len(), 64); + + // Tampering breaks the chain. + let tampered = ndjson.replace("\"y\"", "\"YY\""); + assert!(!verify_hash_chain(&tampered)); + + let parsed = parse_ndjson(&ndjson).unwrap(); + assert_eq!(parsed.len(), 3); + assert_eq!(parsed[2].data["tool_name"], "Bash"); + } + + #[test] + fn meta_round_trips() { + let mut meta = SessionMeta::new("sid".into(), 100.0, "claude-code", "imported".into()); + meta.tool_calls = 6; + meta.total_tokens = 42; + meta.ended_at = Some(200.0); + let text = meta.to_json(); + let back = SessionMeta::from_json(&text).unwrap(); + assert_eq!(back.session_id, "sid"); + assert_eq!(back.tool_calls, 6); + assert_eq!(back.ended_at, Some(200.0)); + // attribution always present, zero fields dropped. + assert!(text.contains("\"attribution\"")); + assert!(!text.contains("\"errors\"")); + } +} From 37eccb92aea3bef0b15b9a659cb16eed46c87505 Mon Sep 17 00:00:00 2001 From: Vishwanath Seshagiri Date: Wed, 17 Jun 2026 15:00:26 -0700 Subject: [PATCH 2/7] Derive event_id from source uuid; add cross-language tests event_id was a per-session counter, so ids collided across sessions and broke consumers that treat them as global (e.g. OTLP span ids). Derive it from each JSONL entry's uuid (12 hex), with an intra-entry suffix and a session_end fallback. Add tests/test_rust_core_roundtrip.py: Python<->Rust hash-chain verification in both directions and byte-identical serialization. --- crates/agent-trace-core/README.md | 15 ++- crates/agent-trace-core/src/import.rs | 102 ++++++++++++++---- tests/test_rust_core_roundtrip.py | 148 ++++++++++++++++++++++++++ 3 files changed, 238 insertions(+), 27 deletions(-) create mode 100644 tests/test_rust_core_roundtrip.py diff --git a/crates/agent-trace-core/README.md b/crates/agent-trace-core/README.md index e932664..c67d2cf 100644 --- a/crates/agent-trace-core/README.md +++ b/crates/agent-trace-core/README.md @@ -58,11 +58,16 @@ The serialized output matches `agent_trace.models` field-for-field (compact event lines, conditional key dropping, pretty meta), so files written here are read by the existing Python tooling and round-trip cleanly. -One intentional difference: imported events get sequential `event_id`s -(`000000000000`, `000000000001`, …) instead of Python's random UUID fragments. -`event_id` isn't used to correlate imported events, and sequential ids make the -hash chain reproducible for tests. The chain itself is always internally -consistent (`verify_hash_chain` passes). +`event_id`s are derived from Claude Code's per-entry `uuid` (dashes stripped, +truncated to 12 hex — the same width as Python's `uuid4().hex[:12]`). This keeps +them stable across re-imports and **globally unique across sessions**, which +matters because the export layers key off `event_id` (e.g. OTLP derives span +ids from it; colliding ids across sessions would corrupt a multi-session +export). The rare entry that yields more than one event gets an intra-entry +suffix; synthetic events with no source uuid (the `session_end` event) derive +their id from the session id. The values won't match a Python import's *random* +ids, but they share the format and uniqueness guarantees, and the hash chain is +always internally consistent (`verify_hash_chain` passes). ## Tests diff --git a/crates/agent-trace-core/src/import.rs b/crates/agent-trace-core/src/import.rs index bccc792..fbae7a6 100644 --- a/crates/agent-trace-core/src/import.rs +++ b/crates/agent-trace-core/src/import.rs @@ -143,6 +143,28 @@ fn usage_u64(usage: &Value, key: &str) -> u64 { usage.get(key).and_then(Value::as_u64).unwrap_or(0) } +/// First 12 hex chars of SHA-256(`s`). +fn hex12(s: &str) -> String { + sha256_hex(s).chars().take(12).collect() +} + +/// Derive a stable, globally-unique 12-hex `event_id` from an entry's `uuid`, +/// falling back to `fallback_seed` when there is none. `intra` disambiguates an +/// entry that yields more than one event. +fn event_id(uuid: Option<&str>, fallback_seed: &str, intra: usize) -> String { + match uuid { + Some(u) if !u.is_empty() => { + let stripped: String = u.chars().filter(|c| *c != '-').collect(); + if intra == 0 { + stripped.chars().take(12).collect() + } else { + hex12(&format!("{}:{}", stripped, intra)) + } + } + _ => hex12(&format!("{}:{}", fallback_seed, intra)), + } +} + /// Import a Claude Code JSONL session log into `//`. pub fn import_jsonl(path: &str, trace_dir: &str) -> io::Result { let path = expanduser(path); @@ -208,29 +230,27 @@ pub fn import_jsonl(path: &str, trace_dir: &str) -> io::Result { ), ); - // Second pass: convert entries to events. event_id uses a per-import - // counter (unique within the session); see README for why this differs - // from the Python random ids. + // Second pass: convert entries to events. let mut events: Vec = Vec::new(); - let mut counter: u64 = 0; - let new_id = |c: &mut u64| -> String { - let id = format!("{:012x}", *c); - *c += 1; - id - }; - for raw in &entries { + for (idx, raw) in entries.iter().enumerate() { let entry_type = obj_get_str(raw, "type"); let ts = parse_iso_timestamp(obj_get_str(raw, "timestamp")); - // Match Python's `msg = raw.get("message", {})`: a missing message is an - // empty object (entry still processed, e.g. `system`/turn_duration); a - // present-but-non-object message skips the entry. let empty = json!({}); let msg = match raw.get("message") { Some(m) if m.is_object() => m, Some(_) => continue, None => &empty, }; + + let entry_uuid = raw.get("uuid").and_then(Value::as_str); + let fallback_seed = format!("{}:{}", session_id, idx); + let mut intra = 0usize; + let mut new_id = || { + let id = event_id(entry_uuid, &fallback_seed, intra); + intra += 1; + id + }; let content = msg.get("content").cloned().unwrap_or(Value::Null); let usage = msg.get("usage").cloned().unwrap_or(Value::Null); let is_sidechain = raw.get("isSidechain").and_then(Value::as_bool).unwrap_or(false); @@ -251,7 +271,7 @@ pub fn import_jsonl(path: &str, trace_dir: &str) -> io::Result { events.push(TraceEvent::new( TOOL_RESULT, ts, - new_id(&mut counter), + new_id(), session_id.clone(), json!({ "tool_use_id": tr.tool_use_id, "content_preview": preview }), )); @@ -270,7 +290,7 @@ pub fn import_jsonl(path: &str, trace_dir: &str) -> io::Result { events.push(TraceEvent::new( TOOL_RESULT, ts, - new_id(&mut counter), + new_id(), session_id.clone(), json!({ "result": result_text, "content_types": ["text"] }), )); @@ -284,7 +304,7 @@ pub fn import_jsonl(path: &str, trace_dir: &str) -> io::Result { events.push(TraceEvent::new( USER_PROMPT, ts, - new_id(&mut counter), + new_id(), session_id.clone(), json!({ "prompt": take_chars(&text_val, 2000) }), )); @@ -314,7 +334,7 @@ pub fn import_jsonl(path: &str, trace_dir: &str) -> io::Result { events.push(TraceEvent::new( TOOL_CALL, ts, - new_id(&mut counter), + new_id(), session_id.clone(), Value::Object(data), )); @@ -326,9 +346,9 @@ pub fn import_jsonl(path: &str, trace_dir: &str) -> io::Result { events.push(TraceEvent::new( ASSISTANT_RESPONSE, ts, - new_id(&mut counter), + new_id(), session_id.clone(), - json!({ "text": take_chars(&text_val, 2000), "model": model }), + json!({ "text": take_chars(&text_val, 2000), "model": obj_get_str(msg, "model") }), )); } @@ -349,7 +369,6 @@ pub fn import_jsonl(path: &str, trace_dir: &str) -> io::Result { } } - // Finalize. meta.ended_at = Some(if last_ts > 0.0 { last_ts } else { meta.started_at }); let ended = meta.ended_at.unwrap(); if meta.total_duration_ms == 0.0 && ended > meta.started_at { @@ -359,7 +378,7 @@ pub fn import_jsonl(path: &str, trace_dir: &str) -> io::Result { events.push(TraceEvent::new( SESSION_END, ended, - new_id(&mut counter), + event_id(None, &format!("{}:session_end", session_id), 0), session_id.clone(), json!({ "duration_ms": compact_number(meta.total_duration_ms), @@ -398,7 +417,46 @@ pub fn import_jsonl(path: &str, trace_dir: &str) -> io::Result { }) } -/// Decode Claude Code's encoded project directory name (matches Python). +#[cfg(test)] +mod tests { + use super::event_id; + + #[test] + fn event_id_reuses_source_uuid() { + let id = event_id(Some("84935712-19fa-4fa0-bf29-3c4341716582"), "seed", 0); + assert_eq!(id, "8493571219fa"); + assert_eq!(id.len(), 12); + } + + #[test] + fn event_id_disambiguates_multiple_events_per_entry() { + let u = Some("84935712-19fa-4fa0-bf29-3c4341716582"); + let a = event_id(u, "seed", 0); + let b = event_id(u, "seed", 1); + let c = event_id(u, "seed", 2); + assert_ne!(a, b); + assert_ne!(b, c); + assert_eq!(b.len(), 12); + } + + #[test] + fn event_id_is_unique_across_sessions() { + let a = event_id(Some("aaaaaaaa-0000-0000-0000-000000000000"), "s1", 0); + let b = event_id(Some("bbbbbbbb-0000-0000-0000-000000000000"), "s2", 0); + assert_ne!(a, b); + } + + #[test] + fn event_id_falls_back_without_uuid() { + let a = event_id(None, "sessionA:session_end", 0); + let b = event_id(None, "sessionB:session_end", 0); + assert_eq!(a.len(), 12); + assert_ne!(a, b); + assert_eq!(a, event_id(None, "sessionA:session_end", 0)); + } +} + +/// Decode Claude Code's encoded project directory name. fn decode_project_path(encoded: &str) -> String { if !encoded.starts_with('-') { return encoded.to_string(); diff --git a/tests/test_rust_core_roundtrip.py b/tests/test_rust_core_roundtrip.py new file mode 100644 index 0000000..6e67321 --- /dev/null +++ b/tests/test_rust_core_roundtrip.py @@ -0,0 +1,148 @@ +"""Cross-language round-trip tests for the Rust core (agent_trace_core). + +Skipped when the Rust wheel isn't installed. +""" + +import hashlib +import json +import tempfile +import unittest +from pathlib import Path + +from agent_trace.models import EventType, TraceEvent +from agent_trace.store import TraceStore + +try: + import agent_trace_core as core + _HAVE_CORE = True +except ImportError: + _HAVE_CORE = False + + +def _python_chain_ok(text: str) -> bool: + """Replicate the Python hash-chain rule for verifying Rust output.""" + prev_line = "" + for line in text.splitlines(): + line = line.strip() + if not line: + continue + prev_hash = json.loads(line).get("prev_hash", "") + expected = hashlib.sha256(prev_line.encode()).hexdigest() if prev_line else "" + if prev_hash != expected: + return False + prev_line = line + return True + + +@unittest.skipUnless(_HAVE_CORE, "agent_trace_core (Rust wheel) not installed") +class RustCoreRoundTrip(unittest.TestCase): + def test_python_writes_rust_verifies(self): + """A hash chain written by the Python store verifies in Rust.""" + with tempfile.TemporaryDirectory() as d: + store = TraceStore(d, redact=False) + from agent_trace.models import SessionMeta + + meta = SessionMeta(session_id="sess", started_at=1.0) + store.create_session(meta) + for i in range(5): + store.append_event( + "sess", + TraceEvent( + event_type=EventType.TOOL_CALL, + timestamp=float(i), + event_id=f"evt{i:09d}", + session_id="sess", + data={"tool_name": "Bash", "n": i}, + ), + ) + text = (Path(d) / "sess" / "events.ndjson").read_text() + + self.assertTrue(core.verify_hash_chain(text)) + tampered = text.replace('"Bash"', '"Sh"', 1) + self.assertFalse(core.verify_hash_chain(tampered)) + + def test_rust_writes_python_verifies(self): + """A session written by Rust passes the Python chain rule and parses.""" + entries = [ + { + "type": "user", + "uuid": "11111111-1111-1111-1111-111111111111", + "sessionId": "rsess", + "timestamp": "2026-06-17T00:00:00.000Z", + "message": {"role": "user", "content": "hello there"}, + }, + { + "type": "assistant", + "uuid": "22222222-2222-2222-2222-222222222222", + "sessionId": "rsess", + "timestamp": "2026-06-17T00:00:01.000Z", + "message": { + "role": "assistant", + "model": "claude-opus-4-8", + "content": [ + {"type": "text", "text": "let me check"}, + {"type": "tool_use", "id": "tu_1", "name": "Bash", "input": {"cmd": "ls"}}, + ], + "usage": {"input_tokens": 10, "output_tokens": 5}, + }, + }, + ] + with tempfile.TemporaryDirectory() as d: + jsonl = Path(d) / "rsess.jsonl" + jsonl.write_text("\n".join(json.dumps(e) for e in entries) + "\n") + trace_dir = Path(d) / "traces" + + summary = core.import_claude_jsonl(str(jsonl), str(trace_dir)) + self.assertEqual(summary["session_id"], "rsess") + + events_file = trace_dir / "rsess" / "events.ndjson" + text = events_file.read_text() + + # Python agrees the Rust-written chain is intact... + self.assertTrue(_python_chain_ok(text)) + # ...and the Python store parses every line into a TraceEvent. + evs = TraceStore(str(trace_dir)).load_events("rsess") + self.assertTrue(len(evs) >= 3) # user_prompt, tool_call, assistant_response, session_end + # event_id derives from the source uuid (dashes stripped, 12 hex). + ids = {e.event_id for e in evs} + self.assertIn("111111111111", ids) + self.assertIn("222222222222", ids) + + def test_serialization_is_byte_identical(self): + """parse_ndjson round-trips Python-serialized lines unchanged — this is + the property the hash chain's cross-language validity depends on. In + particular it catches the `redacted: false` omit-vs-emit divergence.""" + cases = [ + TraceEvent( + event_type=EventType.USER_PROMPT, + timestamp=1.5, + event_id="aaaaaaaaaaaa", + session_id="s", + data={"prompt": "hi"}, + ), + TraceEvent( + event_type=EventType.SESSION_END, + timestamp=0.0, + event_id="bbbbbbbbbbbb", + session_id="", + data={}, + ), + TraceEvent( + event_type=EventType.TOOL_RESULT, + timestamp=2.0, + event_id="cccccccccccc", + session_id="s", + parent_id="aaaaaaaaaaaa", + duration_ms=12.0, + data={"result": "ok"}, + redacted=True, + ), + ] + for ev in cases: + line = ev.to_json() + self.assertEqual(core.parse_ndjson(line), [line], f"mismatch for {line}") + self.assertNotIn("redacted", cases[1].to_json()) + + +if __name__ == "__main__": + unittest.main() From 2958898b1640221ebf83900bc9159a3a4c073b07 Mon Sep 17 00:00:00 2001 From: Vishwanath Seshagiri Date: Wed, 17 Jun 2026 15:26:05 -0700 Subject: [PATCH 3/7] Make PyO3 optional; rebuild hash chain on write; verify verbatim - PyO3 is optional behind a `python` feature; the default rlib build is pure Rust and links no libpython. - write_ndjson recomputes the whole prev_hash chain, discarding stale values; import reuses it instead of duplicating the loop. - verify_hash_chain hashes lines verbatim, matching the writer exactly. --- crates/agent-trace-core/Cargo.toml | 13 ++++++--- crates/agent-trace-core/README.md | 33 ++++++++++++++++++----- crates/agent-trace-core/src/import.rs | 14 +--------- crates/agent-trace-core/src/lib.rs | 15 ++++++----- crates/agent-trace-core/src/models.rs | 39 +++++++++++++++++++-------- 5 files changed, 73 insertions(+), 41 deletions(-) diff --git a/crates/agent-trace-core/Cargo.toml b/crates/agent-trace-core/Cargo.toml index 4576190..d74c0e0 100644 --- a/crates/agent-trace-core/Cargo.toml +++ b/crates/agent-trace-core/Cargo.toml @@ -19,9 +19,14 @@ serde = { version = "1", features = ["derive"] } serde_json = { version = "1", features = ["preserve_order"] } sha2 = "0.10" chrono = { version = "0.4", default-features = false, features = ["std"] } -pyo3 = { version = "0.24", features = ["abi3-py310"] } +# Optional: only pulled in for the Python bindings. Pure-Rust consumers of the +# rlib get no PyO3 and link no libpython. +pyo3 = { version = "0.24", optional = true } [features] -# maturin enables this when building the wheel. Leaving it off lets `cargo test` -# link libpython normally so the pure-Rust unit tests run without a wheel. -extension-module = ["pyo3/extension-module"] +# No default features -> the rlib and `cargo test` build pure Rust, no libpython. +default = [] +# `python` compiles the PyO3 bindings (lib.rs is gated on it). +python = ["dep:pyo3", "pyo3/abi3-py310"] +# maturin enables this when building the wheel; it implies `python`. +extension-module = ["python", "pyo3/extension-module"] diff --git a/crates/agent-trace-core/README.md b/crates/agent-trace-core/README.md index c67d2cf..9074c09 100644 --- a/crates/agent-trace-core/README.md +++ b/crates/agent-trace-core/README.md @@ -3,11 +3,27 @@ The Rust core for [agent-trace](../../README.md). One crate, two surfaces: - **Rust library** (`rlib`) — `models` (trace types, NDJSON, SHA-256 hash chain) - and `import` (Claude Code JSONL import). Use it from any Rust program. + and `import` (Claude Code JSONL import). Use it from any Rust program. The + default build is **pure Rust**: PyO3 is an optional dependency, so a plain Rust + consumer pulls in no PyO3 and links no libpython. - **Python extension** (`cdylib`, built with [maturin](https://www.maturin.rs/) + [PyO3](https://pyo3.rs/)) — the module `agent_trace_core`, exposing the same functionality to Python with no per-event Python object allocation. Parsing - and verifying large traces stays flat in memory. + and verifying large traces stays flat in memory. The bindings compile only + under the `python` feature (which `extension-module` implies). + +### Feature flags + +| Feature | Effect | +|---|---| +| *(default)* | Pure-Rust `rlib`; no PyO3, no libpython. | +| `python` | Compiles the PyO3 bindings (abi3, py3.10+). | +| `extension-module` | `python` + `pyo3/extension-module`; what maturin builds. | + +```toml +# Rust-only consumer — nothing Python is linked: +agent-trace-core = "0.1" +``` ## Why Rust @@ -71,11 +87,16 @@ always internally consistent (`verify_hash_chain` passes). ## Tests -The unit tests are pure Rust, but the test binary links PyO3, so the loader -needs `libpython` on its path: +The default test build is pure Rust, so the core tests run with no libpython: + +```bash +cargo test # models + import, no Python needed +``` + +Building/testing the PyO3 bindings needs `libpython` on the loader path: ```bash LIBDIR=$(python3 -c "import sysconfig; print(sysconfig.get_config_var('LIBDIR'))") -DYLD_FALLBACK_LIBRARY_PATH="$LIBDIR" cargo test # macOS -LD_LIBRARY_PATH="$LIBDIR" cargo test # Linux +DYLD_FALLBACK_LIBRARY_PATH="$LIBDIR" cargo test --features python # macOS +LD_LIBRARY_PATH="$LIBDIR" cargo test --features python # Linux ``` diff --git a/crates/agent-trace-core/src/import.rs b/crates/agent-trace-core/src/import.rs index fbae7a6..e3d1c39 100644 --- a/crates/agent-trace-core/src/import.rs +++ b/crates/agent-trace-core/src/import.rs @@ -389,23 +389,11 @@ pub fn import_jsonl(path: &str, trace_dir: &str) -> io::Result { }), )); - // Write the session directory. let session_dir = Path::new(trace_dir).join(&session_id); fs::create_dir_all(&session_dir)?; fs::write(session_dir.join("meta.json"), meta.to_json())?; - // Build the events file with the SHA-256 hash chain. - let mut ndjson = String::new(); - let mut prev_line = String::new(); - for ev in events.iter_mut() { - if !prev_line.is_empty() { - ev.prev_hash = sha256_hex(&prev_line); - } - let line = ev.to_json(); - ndjson.push_str(&line); - ndjson.push('\n'); - prev_line = line; - } + let ndjson = crate::models::write_ndjson(&mut events); fs::write(session_dir.join("events.ndjson"), ndjson)?; Ok(ImportSummary { diff --git a/crates/agent-trace-core/src/lib.rs b/crates/agent-trace-core/src/lib.rs index ccd057c..6905068 100644 --- a/crates/agent-trace-core/src/lib.rs +++ b/crates/agent-trace-core/src/lib.rs @@ -1,16 +1,16 @@ //! `agent_trace_core` — the Rust core for agent-trace. //! -//! Two surfaces from one crate: -//! * a plain Rust library (`rlib`) — see [`models`] and [`import`]; -//! * a Python extension module (`cdylib`, built by maturin) exposing the -//! same functionality with no per-event Python allocations, so importing -//! and verifying large traces stays flat in memory. +//! Builds as a pure-Rust library (`rlib`) and, under the `python` feature, as a +//! PyO3 extension module (`cdylib`). pub mod import; pub mod models; -use pyo3::prelude::*; -use pyo3::types::{PyDict, PyList}; +#[cfg(feature = "python")] +mod python_bindings { + use crate::{import, models}; + use pyo3::prelude::*; + use pyo3::types::{PyDict, PyList}; /// Import a Claude Code JSONL session log into `//`. /// Returns a summary dict: session_id, tool_calls, llm_requests, total_tokens, events. @@ -72,3 +72,4 @@ fn agent_trace_core(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(verify_hash_chain, m)?)?; Ok(()) } +} // mod python_bindings diff --git a/crates/agent-trace-core/src/models.rs b/crates/agent-trace-core/src/models.rs index a8342e1..b6a3583 100644 --- a/crates/agent-trace-core/src/models.rs +++ b/crates/agent-trace-core/src/models.rs @@ -178,20 +178,21 @@ pub fn sha256_hex(s: &str) -> String { out } -/// Serialize a list of events into an NDJSON string, filling each event's -/// `prev_hash` with the SHA-256 of the previous line (empty for the first). -/// The events are mutated in place so the returned chain is self-consistent. +/// Serialize events to NDJSON, rebuilding the `prev_hash` chain (first empty, +/// each subsequent one the SHA-256 of the previous line). Events are mutated in +/// place to reflect what was written. pub fn write_ndjson(events: &mut [TraceEvent]) -> String { let mut out = String::new(); - let mut prev_line = String::new(); + let mut prev_line: Option = None; for ev in events.iter_mut() { - if ev.prev_hash.is_empty() && !prev_line.is_empty() { - ev.prev_hash = sha256_hex(&prev_line); - } + ev.prev_hash = match &prev_line { + Some(prev) => sha256_hex(prev), + None => String::new(), + }; let line = ev.to_json(); out.push_str(&line); out.push('\n'); - prev_line = line; + prev_line = Some(line); } out } @@ -209,12 +210,10 @@ pub fn parse_ndjson(text: &str) -> serde_json::Result> { Ok(events) } -/// Verify that every line's `prev_hash` equals the SHA-256 of the previous -/// line. Returns true for an intact chain (and for empty input). +/// Verify that every line's `prev_hash` equals the SHA-256 of the previous line. pub fn verify_hash_chain(text: &str) -> bool { let mut prev_line: Option<&str> = None; for line in text.lines() { - let line = line.trim(); if line.is_empty() { continue; } @@ -280,6 +279,24 @@ mod tests { assert_eq!(parsed[2].data["tool_name"], "Bash"); } + #[test] + fn write_ndjson_overwrites_stale_prev_hash() { + // Events arrive carrying bogus prev_hash values; write_ndjson must + // discard them and rebuild a consistent chain. + let mut events = vec![ + TraceEvent::new(USER_PROMPT, 1.0, "a".into(), "s".into(), json!({"prompt": "x"})), + TraceEvent::new(ASSISTANT_RESPONSE, 2.0, "b".into(), "s".into(), json!({"text": "y"})), + ]; + events[0].prev_hash = "deadbeef".into(); // stale on first + events[1].prev_hash = "garbage".into(); // stale mid-chain + let ndjson = write_ndjson(&mut events); + assert!(verify_hash_chain(&ndjson)); + // Both stale values are discarded: first reset to empty, second recomputed. + assert_eq!(events[0].prev_hash, ""); + assert_eq!(events[1].prev_hash.len(), 64); + assert_ne!(events[1].prev_hash, "garbage"); + } + #[test] fn meta_round_trips() { let mut meta = SessionMeta::new("sid".into(), 100.0, "claude-code", "imported".into()); From eb2c6be6af242947e8a4c93a5124e46de3c80a0e Mon Sep 17 00:00:00 2001 From: Vishwanath Seshagiri Date: Wed, 17 Jun 2026 20:37:38 -0700 Subject: [PATCH 4/7] Escape non-ASCII to match Python's ensure_ascii serde_json emits raw UTF-8 in strings; Python escapes non-ASCII as \uXXXX, so traces with accents/CJK/emoji serialized to different bytes and broke hash-chain parity. to_json now escapes non-ASCII (surrogate pairs for astral chars); ASCII output takes a fast path. --- crates/agent-trace-core/src/models.rs | 43 ++++++++++++++++++-- tests/test_rust_core_roundtrip.py | 57 +++++++++++++++++++++++---- 2 files changed, 89 insertions(+), 11 deletions(-) diff --git a/crates/agent-trace-core/src/models.rs b/crates/agent-trace-core/src/models.rs index b6a3583..369f26c 100644 --- a/crates/agent-trace-core/src/models.rs +++ b/crates/agent-trace-core/src/models.rs @@ -14,6 +14,25 @@ fn is_false(b: &bool) -> bool { !*b } +/// Escape non-ASCII as `\uXXXX` (UTF-16 units), matching Python's `ensure_ascii`. +fn escape_non_ascii(s: String) -> String { + if s.is_ascii() { + return s; + } + let mut out = String::with_capacity(s.len() + 8); + let mut buf = [0u16; 2]; + for ch in s.chars() { + if ch.is_ascii() { + out.push(ch); + } else { + for unit in ch.encode_utf16(&mut buf) { + out.push_str(&format!("\\u{:04x}", unit)); + } + } + } + out +} + fn empty_object() -> Value { Value::Object(serde_json::Map::new()) } @@ -53,10 +72,9 @@ impl TraceEvent { } } - /// Compact JSON, byte-for-byte equivalent to the Python `to_json` output - /// (`separators=(",", ":")`, conditional key dropping, field ordering). + /// Serialize to a compact JSON line. pub fn to_json(&self) -> String { - serde_json::to_string(self).expect("TraceEvent serializes") + escape_non_ascii(serde_json::to_string(self).expect("TraceEvent serializes")) } pub fn from_json(line: &str) -> serde_json::Result { @@ -158,7 +176,7 @@ impl SessionMeta { } pub fn to_json(&self) -> String { - serde_json::to_string_pretty(self).expect("SessionMeta serializes") + escape_non_ascii(serde_json::to_string_pretty(self).expect("SessionMeta serializes")) } pub fn from_json(text: &str) -> serde_json::Result { @@ -297,6 +315,23 @@ mod tests { assert_ne!(events[1].prev_hash, "garbage"); } + #[test] + fn non_ascii_is_escaped_like_python() { + let original = "café 日本語 🎉 ❤"; + let ev = TraceEvent::new(ASSISTANT_RESPONSE, 1.0, "abc".into(), "s".into(), json!({"text": original})); + let line = ev.to_json(); + assert!(line.is_ascii(), "output must be pure ASCII, got: {line}"); + let u = |c: char| { + let mut b = [0u16; 2]; + c.encode_utf16(&mut b).iter().map(|x| format!("\\u{:04x}", x)).collect::() + }; + assert!(line.contains(&u('é'))); + assert!(line.contains(&u('🎉'))); + assert_eq!(u('🎉').matches("\\u").count(), 2); + let parsed = TraceEvent::from_json(&line).unwrap(); + assert_eq!(parsed.data["text"], original); + } + #[test] fn meta_round_trips() { let mut meta = SessionMeta::new("sid".into(), 100.0, "claude-code", "imported".into()); diff --git a/tests/test_rust_core_roundtrip.py b/tests/test_rust_core_roundtrip.py index 6e67321..00b7015 100644 --- a/tests/test_rust_core_roundtrip.py +++ b/tests/test_rust_core_roundtrip.py @@ -98,20 +98,63 @@ def test_rust_writes_python_verifies(self): events_file = trace_dir / "rsess" / "events.ndjson" text = events_file.read_text() - # Python agrees the Rust-written chain is intact... self.assertTrue(_python_chain_ok(text)) - # ...and the Python store parses every line into a TraceEvent. evs = TraceStore(str(trace_dir)).load_events("rsess") - self.assertTrue(len(evs) >= 3) # user_prompt, tool_call, assistant_response, session_end - # event_id derives from the source uuid (dashes stripped, 12 hex). + self.assertTrue(len(evs) >= 3) ids = {e.event_id for e in evs} self.assertIn("111111111111", ids) self.assertIn("222222222222", ids) + def test_non_ascii_serialization_matches_python(self): + """Non-ASCII serializes to the same bytes as Python (ensure_ascii).""" + ev = TraceEvent( + event_type=EventType.ASSISTANT_RESPONSE, + timestamp=1.0, + event_id="abc", + session_id="s", + data={"text": "café 日本語 🎉 ❤ \U0001F600"}, + ) + line = ev.to_json() + self.assertTrue(line.isascii()) + self.assertEqual(core.parse_ndjson(line), [line]) + + from agent_trace.jsonl_import import import_jsonl + + entries = [ + { + "type": "user", + "uuid": "33333333-3333-3333-3333-333333333333", + "sessionId": "usess", + "timestamp": "2026-06-17T00:00:00.000Z", + "message": {"role": "user", "content": "héllo 世界 🚀"}, + } + ] + with tempfile.TemporaryDirectory() as d: + jsonl = Path(d) / "usess.jsonl" + jsonl.write_text( + "\n".join(json.dumps(e) for e in entries) + "\n", encoding="utf-8" + ) + py_dir, rust_dir = Path(d) / "py", Path(d) / "rust" + import_jsonl(str(jsonl), store=TraceStore(str(py_dir), redact=False)) + core.import_claude_jsonl(str(jsonl), str(rust_dir)) + + def norm(p): + out = [] + for ln in Path(p).read_text().splitlines(): + if not ln.strip(): + continue + obj = json.loads(ln) + obj.pop("event_id", None) + obj.pop("prev_hash", None) + out.append(json.dumps(obj, sort_keys=True, ensure_ascii=True)) + return out + + py = norm(py_dir / "usess" / "events.ndjson") + rust = norm(rust_dir / "usess" / "events.ndjson") + self.assertEqual(py, rust) + def test_serialization_is_byte_identical(self): - """parse_ndjson round-trips Python-serialized lines unchanged — this is - the property the hash chain's cross-language validity depends on. In - particular it catches the `redacted: false` omit-vs-emit divergence.""" + """parse_ndjson round-trips Python-serialized lines unchanged.""" cases = [ TraceEvent( event_type=EventType.USER_PROMPT, From 7cae28b5722169f4e335be63f6749ed2b4e23b2e Mon Sep 17 00:00:00 2001 From: Vishwanath Seshagiri Date: Fri, 19 Jun 2026 10:39:26 -0700 Subject: [PATCH 5/7] made changes to the parity --- crates/agent-trace-core/README.md | 2 + docs/commands.md | 2 +- docs/setup.md | 2 + src/agent_trace/__init__.py | 2 +- src/agent_trace/jsonl_import.py | 32 ++++++ tests/test_jsonl_import.py | 177 ++++++++++++++++++++++++++++++ tests/test_rust_core_roundtrip.py | 4 +- 7 files changed, 218 insertions(+), 3 deletions(-) diff --git a/crates/agent-trace-core/README.md b/crates/agent-trace-core/README.md index 9074c09..6d17fee 100644 --- a/crates/agent-trace-core/README.md +++ b/crates/agent-trace-core/README.md @@ -70,6 +70,8 @@ let intact = verify_hash_chain(&std::fs::read_to_string("events.ndjson")?); ## Compatibility with the Python implementation +The Python package treats this extension as an optional accelerator for Claude JSONL import. `agent_trace.jsonl_import.import_jsonl()` uses it only when the module is installed, redaction is disabled, no workspace is active, and `AGENT_STRACE_NO_RUST` is unset; otherwise the standard-library Python importer remains the fallback. + The serialized output matches `agent_trace.models` field-for-field (compact event lines, conditional key dropping, pretty meta), so files written here are read by the existing Python tooling and round-trip cleanly. diff --git a/docs/commands.md b/docs/commands.md index 6456b71..84f5c60 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -38,7 +38,7 @@ Print or install hooks config for supported agent CLIs. `--cli claude` prints Cl ``` agent-strace import [--discover] ``` -Import a Claude Code JSONL session log. `--discover` lists available sessions in `~/.claude/projects/`. +Import a Claude Code JSONL session log. `--discover` lists available sessions in `~/.claude/projects/`. When the optional `agent_trace_core` extension is installed, imports use the Rust path only with redaction disabled and no active workspace; `AGENT_STRACE_NO_RUST=1` forces the Python fallback. --- diff --git a/docs/setup.md b/docs/setup.md index dc39136..b62b39a 100644 --- a/docs/setup.md +++ b/docs/setup.md @@ -182,6 +182,8 @@ agent-strace import --discover agent-strace import ~/.claude/projects//.jsonl ``` +Import is always available through the Python parser. If the optional `agent_trace_core` Rust extension is installed, `agent-strace import` uses it only when redaction is disabled and no workspace is active; set `AGENT_STRACE_NO_RUST=1` to force the Python path. + --- ## Option 2: MCP proxy (any MCP client) diff --git a/src/agent_trace/__init__.py b/src/agent_trace/__init__.py index 8bf457d..525742d 100644 --- a/src/agent_trace/__init__.py +++ b/src/agent_trace/__init__.py @@ -1,3 +1,3 @@ """agent-trace: strace for AI agents.""" -__version__ = "0.80.1" +__version__ = "0.80.2" diff --git a/src/agent_trace/jsonl_import.py b/src/agent_trace/jsonl_import.py index 5aa4dd7..6b8f90b 100755 --- a/src/agent_trace/jsonl_import.py +++ b/src/agent_trace/jsonl_import.py @@ -15,6 +15,7 @@ import argparse import json +import os import sys from pathlib import Path from typing import Any @@ -136,6 +137,33 @@ def _extract_text(content: Any) -> str: return "" +def _rust_core() -> Any | None: + """Return the optional Rust extension module when enabled and installed.""" + if os.environ.get("AGENT_STRACE_NO_RUST", "").strip().lower() in ( + "1", + "true", + "yes", + ): + return None + try: + import agent_trace_core + + return agent_trace_core + except ImportError: + return None + + +def _rust_import(path: Path, store: TraceStore) -> str | None: + """Import via the Rust core when eligible; return the session id, or None + to fall back to the Python implementation.""" + if store.redact or store.workspace_id: + return None + core = _rust_core() + if core is None: + return None + return core.import_claude_jsonl(str(path), str(store.base_dir))["session_id"] + + def import_jsonl( path: str | Path, store: TraceStore | None = None, @@ -164,6 +192,10 @@ def import_jsonl( if store is None: store = TraceStore(trace_dir) + rust_session_id = _rust_import(path, store) + if rust_session_id is not None: + return rust_session_id + # First pass: extract metadata from first entry session_id = "" extra_meta: dict[str, str] = {} diff --git a/tests/test_jsonl_import.py b/tests/test_jsonl_import.py index d45031e..496fbd2 100644 --- a/tests/test_jsonl_import.py +++ b/tests/test_jsonl_import.py @@ -5,9 +5,13 @@ import sys import tempfile import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) +import agent_trace.jsonl_import as jsonl_import from agent_trace.jsonl_import import import_jsonl from agent_trace.models import EventType from agent_trace.store import TraceStore @@ -26,6 +30,23 @@ def _write_jsonl(self, entries): f.write(json.dumps(entry) + "\n") return path + def _failing_rust_core(self, message): + def fail(*_args): + raise AssertionError(message) + + return SimpleNamespace(import_claude_jsonl=fail) + + def _single_user_jsonl(self, session_id, content="Hello"): + return self._write_jsonl([ + { + "type": "user", + "sessionId": session_id, + "timestamp": "2025-01-01T00:00:00Z", + "message": {"role": "user", "content": content}, + "uuid": "u1", + }, + ]) + def test_import_basic_session(self): path = self._write_jsonl([ { @@ -168,6 +189,162 @@ def test_import_nonexistent_file(self): with self.assertRaises(FileNotFoundError): import_jsonl("/nonexistent/path.jsonl", store=self.store) + def test_falls_back_when_rust_core_absent(self): + path = self._single_user_jsonl("fallback123") + store = TraceStore(os.path.join(self.tmpdir, "fallback"), redact=False) + + with patch.object(jsonl_import, "_rust_core", return_value=None) as rust_core: + session_id = import_jsonl(path, store=store) + + rust_core.assert_called_once() + self.assertEqual(session_id, "fallback123") + self.assertEqual(store.load_meta(session_id).agent_name, "claude-code") + + def test_redaction_forces_python_path(self): + path = self._single_user_jsonl( + "redact123", "secret sk-abc123def456ghi789jkl012mno345pqr678" + ) + store = TraceStore(os.path.join(self.tmpdir, "redact"), redact=True) + with patch.object( + jsonl_import, + "_rust_core", + side_effect=AssertionError("Rust helper should not run with redaction"), + ): + session_id = import_jsonl(path, store=store) + + events = store.load_events(session_id) + prompts = [e for e in events if e.event_type == EventType.USER_PROMPT] + self.assertEqual(len(prompts), 1) + self.assertIn("[REDACTED:openai-key]", prompts[0].data["prompt"]) + + def test_workspace_forces_python_path(self): + path = self._single_user_jsonl("workspace123") + store = TraceStore( + os.path.join(self.tmpdir, "workspace"), workspace_id="ws1", redact=False + ) + with patch.object( + jsonl_import, + "_rust_core", + side_effect=AssertionError("Rust helper should not run for workspaces"), + ): + session_id = import_jsonl(path, store=store) + + meta = store.load_meta(session_id) + self.assertEqual(meta.workspace_id, "ws1") + + def test_no_rust_env_forces_python_path(self): + path = self._single_user_jsonl("norust123") + store = TraceStore(os.path.join(self.tmpdir, "norust"), redact=False) + rust_core = self._failing_rust_core( + "Rust path should not run when AGENT_STRACE_NO_RUST is set" + ) + + with patch.dict(os.environ, {"AGENT_STRACE_NO_RUST": "1"}): + with patch.dict(sys.modules, {"agent_trace_core": rust_core}): + session_id = import_jsonl(path, store=store) + + self.assertEqual(session_id, "norust123") + self.assertEqual(store.load_meta(session_id).agent_name, "claude-code") + + def test_rust_path_matches_python_content_when_available(self): + try: + import agent_trace_core as core + except ImportError: + self.skipTest("agent_trace_core (Rust wheel) not installed") + + entries = [ + { + "type": "user", + "sessionId": "parity123", + "gitBranch": "main", + "version": "2.1.50", + "timestamp": "2025-01-01T00:00:00.250Z", + "message": {"role": "user", "content": "Hello"}, + "uuid": "11111111-1111-1111-1111-111111111111", + }, + { + "type": "assistant", + "sessionId": "parity123", + "timestamp": "2025-01-01T00:00:01.500Z", + "message": { + "role": "assistant", + "model": "claude-sonnet-4-20250514", + "content": [ + {"type": "text", "text": "I'll inspect it."}, + { + "type": "tool_use", + "id": "tool1", + "name": "Bash", + "input": {"command": "ls -la"}, + "caller": {"type": "direct"}, + }, + ], + "usage": { + "input_tokens": 100, + "output_tokens": 50, + "cache_creation_input_tokens": 2, + "cache_read_input_tokens": 3, + }, + }, + "uuid": "22222222-2222-2222-2222-222222222222", + "parentUuid": "11111111-1111-1111-1111-111111111111", + }, + { + "type": "user", + "sessionId": "parity123", + "timestamp": "2025-01-01T00:00:02.000Z", + "message": { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "tool1", + "content": "file output", + } + ], + }, + "uuid": "33333333-3333-3333-3333-333333333333", + }, + { + "type": "system", + "sessionId": "parity123", + "timestamp": "2025-01-01T00:00:02.750Z", + "subtype": "turn_duration", + "durationMs": 432.5, + }, + ] + + with tempfile.TemporaryDirectory() as d: + jsonl = Path(d) / "session.jsonl" + jsonl.write_text("\n".join(json.dumps(e) for e in entries) + "\n") + py_store = TraceStore(Path(d) / "python", redact=False) + rust_store = TraceStore(Path(d) / "rust", redact=False) + + with patch.object(jsonl_import, "_rust_core", return_value=None): + py_session_id = import_jsonl(jsonl, store=py_store) + with patch.object(jsonl_import, "_rust_core", return_value=core): + rust_session_id = import_jsonl(jsonl, store=rust_store) + + self.assertEqual(py_session_id, rust_session_id) + self.assertEqual( + json.loads(py_store.load_meta(py_session_id).to_json()), + json.loads(rust_store.load_meta(rust_session_id).to_json()), + ) + + def normalized_events(store, session_id): + normalized = [] + for event in store.load_events(session_id): + obj = json.loads(event.to_json()) + obj.pop("event_id", None) + obj.pop("prev_hash", None) + normalized.append(obj) + return normalized + + self.assertEqual( + normalized_events(py_store, py_session_id), + normalized_events(rust_store, rust_session_id), + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_rust_core_roundtrip.py b/tests/test_rust_core_roundtrip.py index 00b7015..5ec21bb 100644 --- a/tests/test_rust_core_roundtrip.py +++ b/tests/test_rust_core_roundtrip.py @@ -8,6 +8,7 @@ import tempfile import unittest from pathlib import Path +from unittest.mock import patch from agent_trace.models import EventType, TraceEvent from agent_trace.store import TraceStore @@ -135,7 +136,8 @@ def test_non_ascii_serialization_matches_python(self): "\n".join(json.dumps(e) for e in entries) + "\n", encoding="utf-8" ) py_dir, rust_dir = Path(d) / "py", Path(d) / "rust" - import_jsonl(str(jsonl), store=TraceStore(str(py_dir), redact=False)) + with patch("agent_trace.jsonl_import._rust_core", return_value=None): + import_jsonl(str(jsonl), store=TraceStore(str(py_dir), redact=False)) core.import_claude_jsonl(str(jsonl), str(rust_dir)) def norm(p): From b4813697a158bc039477287a09c0bc35df145d68 Mon Sep 17 00:00:00 2001 From: Vishwanath Seshagiri Date: Fri, 19 Jun 2026 12:47:35 -0700 Subject: [PATCH 6/7] changed README and imports --- crates/agent-trace-core/README.md | 44 +++++++++----------------- crates/agent-trace-core/src/import.rs | 45 +++++++++++++-------------- 2 files changed, 35 insertions(+), 54 deletions(-) diff --git a/crates/agent-trace-core/README.md b/crates/agent-trace-core/README.md index 6d17fee..623e36a 100644 --- a/crates/agent-trace-core/README.md +++ b/crates/agent-trace-core/README.md @@ -3,13 +3,10 @@ The Rust core for [agent-trace](../../README.md). One crate, two surfaces: - **Rust library** (`rlib`) — `models` (trace types, NDJSON, SHA-256 hash chain) - and `import` (Claude Code JSONL import). Use it from any Rust program. The - default build is **pure Rust**: PyO3 is an optional dependency, so a plain Rust - consumer pulls in no PyO3 and links no libpython. + and `import` (Claude Code JSONL import). Default build is pure Rust; PyO3 is an + optional dependency. - **Python extension** (`cdylib`, built with [maturin](https://www.maturin.rs/) - + [PyO3](https://pyo3.rs/)) — the module `agent_trace_core`, exposing the same - functionality to Python with no per-event Python object allocation. Parsing - and verifying large traces stays flat in memory. The bindings compile only + + [PyO3](https://pyo3.rs/)) — the module `agent_trace_core`. Compiled only under the `python` feature (which `extension-module` implies). ### Feature flags @@ -21,17 +18,10 @@ The Rust core for [agent-trace](../../README.md). One crate, two surfaces: | `extension-module` | `python` + `pyo3/extension-module`; what maturin builds. | ```toml -# Rust-only consumer — nothing Python is linked: +# Rust-only consumer: agent-trace-core = "0.1" ``` -## Why Rust - -The hot paths — parsing NDJSON, walking the hash chain, importing multi-MB -Claude session logs — allocate one short-lived object per event in Python. In -Rust those are stack values dropped deterministically, so memory stays bounded -regardless of trace size, and there's no GC to chase. - ## Build ```bash @@ -70,29 +60,23 @@ let intact = verify_hash_chain(&std::fs::read_to_string("events.ndjson")?); ## Compatibility with the Python implementation -The Python package treats this extension as an optional accelerator for Claude JSONL import. `agent_trace.jsonl_import.import_jsonl()` uses it only when the module is installed, redaction is disabled, no workspace is active, and `AGENT_STRACE_NO_RUST` is unset; otherwise the standard-library Python importer remains the fallback. +`agent_trace.jsonl_import.import_jsonl()` uses this extension when it's +installed, redaction is disabled, no workspace is active, and +`AGENT_STRACE_NO_RUST` is unset; otherwise it falls back to the Python importer. -The serialized output matches `agent_trace.models` field-for-field (compact -event lines, conditional key dropping, pretty meta), so files written here are -read by the existing Python tooling and round-trip cleanly. +Serialized output matches `agent_trace.models` field-for-field: compact event +lines, conditional key dropping, pretty meta, non-ASCII escaped as `\uXXXX`. -`event_id`s are derived from Claude Code's per-entry `uuid` (dashes stripped, -truncated to 12 hex — the same width as Python's `uuid4().hex[:12]`). This keeps -them stable across re-imports and **globally unique across sessions**, which -matters because the export layers key off `event_id` (e.g. OTLP derives span -ids from it; colliding ids across sessions would corrupt a multi-session -export). The rare entry that yields more than one event gets an intra-entry -suffix; synthetic events with no source uuid (the `session_end` event) derive -their id from the session id. The values won't match a Python import's *random* -ids, but they share the format and uniqueness guarantees, and the hash chain is -always internally consistent (`verify_hash_chain` passes). +`event_id`s are derived from each entry's `uuid` (dashes stripped, 12 hex). A +multi-event entry gets an intra-entry suffix; the synthetic `session_end` event +derives its id from the session id. ## Tests -The default test build is pure Rust, so the core tests run with no libpython: +Core tests run on the default (pure-Rust) build, no libpython needed: ```bash -cargo test # models + import, no Python needed +cargo test # models + import ``` Building/testing the PyO3 bindings needs `libpython` on the loader path: diff --git a/crates/agent-trace-core/src/import.rs b/crates/agent-trace-core/src/import.rs index e3d1c39..268d889 100644 --- a/crates/agent-trace-core/src/import.rs +++ b/crates/agent-trace-core/src/import.rs @@ -56,10 +56,6 @@ fn take_chars(s: &str, n: usize) -> String { s.chars().take(n).collect() } -fn char_len(s: &str) -> usize { - s.chars().count() -} - /// Extract text from message content (a string, or a list of content blocks). fn extract_text(content: &Value) -> String { match content { @@ -135,12 +131,12 @@ fn extract_tool_results(content: &[Value]) -> Vec { results } -fn obj_get_str<'a>(v: &'a Value, key: &str) -> &'a str { +fn get_str<'a>(v: &'a Value, key: &str) -> &'a str { v.get(key).and_then(Value::as_str).unwrap_or("") } -fn usage_u64(usage: &Value, key: &str) -> u64 { - usage.get(key).and_then(Value::as_u64).unwrap_or(0) +fn get_u64(v: &Value, key: &str) -> u64 { + v.get(key).and_then(Value::as_u64).unwrap_or(0) } /// First 12 hex chars of SHA-256(`s`). @@ -199,7 +195,7 @@ pub fn import_jsonl(path: &str, trace_dir: &str) -> io::Result { continue; } - let ts = parse_iso_timestamp(obj_get_str(&raw, "timestamp")); + let ts = parse_iso_timestamp(get_str(&raw, "timestamp")); if first_ts == 0.0 && ts > 0.0 { first_ts = ts; } @@ -207,9 +203,9 @@ pub fn import_jsonl(path: &str, trace_dir: &str) -> io::Result { last_ts = ts; } if session_id.is_empty() { - session_id = obj_get_str(&raw, "sessionId").to_string(); - git_branch = obj_get_str(&raw, "gitBranch").to_string(); - version = obj_get_str(&raw, "version").to_string(); + session_id = get_str(&raw, "sessionId").to_string(); + git_branch = get_str(&raw, "gitBranch").to_string(); + version = get_str(&raw, "version").to_string(); } entries.push(raw); } @@ -234,8 +230,8 @@ pub fn import_jsonl(path: &str, trace_dir: &str) -> io::Result { let mut events: Vec = Vec::new(); for (idx, raw) in entries.iter().enumerate() { - let entry_type = obj_get_str(raw, "type"); - let ts = parse_iso_timestamp(obj_get_str(raw, "timestamp")); + let entry_type = get_str(raw, "type"); + let ts = parse_iso_timestamp(get_str(raw, "timestamp")); let empty = json!({}); let msg = match raw.get("message") { Some(m) if m.is_object() => m, @@ -263,7 +259,7 @@ pub fn import_jsonl(path: &str, trace_dir: &str) -> io::Result { for tr in &tool_results { let preview = if tr.content.is_empty() { String::new() - } else if char_len(&tr.content) > 2000 { + } else if tr.content.chars().count() > 2000 { format!("{}...", take_chars(&tr.content, 2000)) } else { tr.content.clone() @@ -280,8 +276,8 @@ pub fn import_jsonl(path: &str, trace_dir: &str) -> io::Result { if tool_results.is_empty() { if let Some(tr_data) = raw.get("toolUseResult") { if tr_data.is_object() { - let stdout = obj_get_str(tr_data, "stdout"); - let stderr = obj_get_str(tr_data, "stderr"); + let stdout = get_str(tr_data, "stdout"); + let stderr = get_str(tr_data, "stderr"); if !stdout.is_empty() || !stderr.is_empty() { let mut result_text = take_chars(stdout, 500); if !stderr.is_empty() { @@ -321,12 +317,12 @@ pub fn import_jsonl(path: &str, trace_dir: &str) -> io::Result { if is_sidechain { data.insert("is_sidechain".into(), json!(true)); } - let caller_type = obj_get_str(&tc.caller, "type"); + let caller_type = get_str(&tc.caller, "type"); if !caller_type.is_empty() { data.insert("caller_type".into(), json!(caller_type)); } if tc.name == "Agent" { - let subagent = obj_get_str(&tc.input, "subagent_type"); + let subagent = get_str(&tc.input, "subagent_type"); if !subagent.is_empty() { data.insert("subagent_type".into(), json!(subagent)); } @@ -348,19 +344,20 @@ pub fn import_jsonl(path: &str, trace_dir: &str) -> io::Result { ts, new_id(), session_id.clone(), - json!({ "text": take_chars(&text_val, 2000), "model": obj_get_str(msg, "model") }), + json!({ "text": take_chars(&text_val, 2000), "model": get_str(msg, "model") }), )); } if usage.is_object() { - meta.total_tokens += usage_u64(&usage, "input_tokens") - + usage_u64(&usage, "output_tokens") - + usage_u64(&usage, "cache_creation_input_tokens") - + usage_u64(&usage, "cache_read_input_tokens"); + // TODO: Add token breakdown later (per-type tokens + model -> cost). + meta.total_tokens += get_u64(&usage, "input_tokens") + + get_u64(&usage, "output_tokens") + + get_u64(&usage, "cache_creation_input_tokens") + + get_u64(&usage, "cache_read_input_tokens"); meta.llm_requests += 1; } } else if entry_type == "system" { - if obj_get_str(raw, "subtype") == "turn_duration" { + if get_str(raw, "subtype") == "turn_duration" { let duration_ms = raw.get("durationMs").and_then(Value::as_f64).unwrap_or(0.0); if duration_ms != 0.0 { meta.total_duration_ms += duration_ms; From 043c1da3965fa2615e91e30582b9f33a45fc0a9b Mon Sep 17 00:00:00 2001 From: vishwanath1306 Date: Mon, 22 Jun 2026 17:43:23 -0700 Subject: [PATCH 7/7] Restore shortened Why Rust section to crate README --- crates/agent-trace-core/README.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/crates/agent-trace-core/README.md b/crates/agent-trace-core/README.md index 623e36a..930bdfe 100644 --- a/crates/agent-trace-core/README.md +++ b/crates/agent-trace-core/README.md @@ -9,6 +9,12 @@ The Rust core for [agent-trace](../../README.md). One crate, two surfaces: + [PyO3](https://pyo3.rs/)) — the module `agent_trace_core`. Compiled only under the `python` feature (which `extension-module` implies). +## Why Rust + +The hot paths allocate one short-lived object per event in Python; in Rust +those are stack values dropped deterministically, so memory stays bounded +regardless of trace size and there's no GC. + ### Feature flags | Feature | Effect |