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..d74c0e0 --- /dev/null +++ b/crates/agent-trace-core/Cargo.toml @@ -0,0 +1,32 @@ +[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"] } +# 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] +# 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 new file mode 100644 index 0000000..930bdfe --- /dev/null +++ b/crates/agent-trace-core/README.md @@ -0,0 +1,94 @@ +# 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). 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`. 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 | +|---|---| +| *(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: +agent-trace-core = "0.1" +``` + +## 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 + +`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. + +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 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 + +Core tests run on the default (pure-Rust) build, no libpython needed: + +```bash +cargo test # models + import +``` + +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 --features python # macOS +LD_LIBRARY_PATH="$LIBDIR" cargo test --features python # 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..268d889 --- /dev/null +++ b/crates/agent-trace-core/src/import.rs @@ -0,0 +1,488 @@ +//! 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() +} + +/// 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 get_str<'a>(v: &'a Value, key: &str) -> &'a str { + v.get(key).and_then(Value::as_str).unwrap_or("") +} + +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`). +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); + 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(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 = get_str(&raw, "sessionId").to_string(); + git_branch = get_str(&raw, "gitBranch").to_string(); + version = 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. + let mut events: Vec = Vec::new(); + + for (idx, raw) in entries.iter().enumerate() { + 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, + 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); + + 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 tr.content.chars().count() > 2000 { + format!("{}...", take_chars(&tr.content, 2000)) + } else { + tr.content.clone() + }; + events.push(TraceEvent::new( + TOOL_RESULT, + ts, + new_id(), + 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 = 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() { + result_text.push_str(&format!(" [stderr: {}]", take_chars(stderr, 200))); + } + events.push(TraceEvent::new( + TOOL_RESULT, + ts, + new_id(), + 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(), + 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 = get_str(&tc.caller, "type"); + if !caller_type.is_empty() { + data.insert("caller_type".into(), json!(caller_type)); + } + if tc.name == "Agent" { + let subagent = 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(), + session_id.clone(), + Value::Object(data), + )); + meta.tool_calls += 1; + } + } + + if !text_val.is_empty() { + events.push(TraceEvent::new( + ASSISTANT_RESPONSE, + ts, + new_id(), + session_id.clone(), + json!({ "text": take_chars(&text_val, 2000), "model": get_str(msg, "model") }), + )); + } + + if usage.is_object() { + // 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 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; + } + } + } + } + + 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, + event_id(None, &format!("{}:session_end", session_id), 0), + 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(), + }), + )); + + 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())?; + + let ndjson = crate::models::write_ndjson(&mut events); + 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(), + }) +} + +#[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(); + } + 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..6905068 --- /dev/null +++ b/crates/agent-trace-core/src/lib.rs @@ -0,0 +1,75 @@ +//! `agent_trace_core` — the Rust core for agent-trace. +//! +//! Builds as a pure-Rust library (`rlib`) and, under the `python` feature, as a +//! PyO3 extension module (`cdylib`). + +pub mod import; +pub mod models; + +#[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. +#[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(()) +} +} // mod python_bindings diff --git a/crates/agent-trace-core/src/models.rs b/crates/agent-trace-core/src/models.rs new file mode 100644 index 0000000..369f26c --- /dev/null +++ b/crates/agent-trace-core/src/models.rs @@ -0,0 +1,350 @@ +//! 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 +} + +/// 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()) +} + +/// 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, + } + } + + /// Serialize to a compact JSON line. + pub fn to_json(&self) -> String { + escape_non_ascii(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 { + escape_non_ascii(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 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: Option = None; + for ev in events.iter_mut() { + 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 = Some(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. +pub fn verify_hash_chain(text: &str) -> bool { + let mut prev_line: Option<&str> = None; + for line in text.lines() { + 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 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 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()); + 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\"")); + } +} 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 new file mode 100644 index 0000000..5ec21bb --- /dev/null +++ b/tests/test_rust_core_roundtrip.py @@ -0,0 +1,193 @@ +"""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 unittest.mock import patch + +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() + + self.assertTrue(_python_chain_ok(text)) + evs = TraceStore(str(trace_dir)).load_events("rsess") + 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" + 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): + 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.""" + 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()