From 7fb7445c9a41197282248bc76f85c35e2c1a0465 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Grabowski?= <> Date: Fri, 3 Apr 2026 20:49:08 +0200 Subject: [PATCH 1/5] feat: add AgentAdapter abstraction with Codex CLI support Replace hardcoded Claude Code transcript parsing with an extensible AgentAdapter trait and registry. Each agent gets its own adapter for event mapping, file change extraction, and transcript record parsing. - AgentAdapter trait with ClaudeCode, Codex, and Default adapters - Codex transcript parsing (response_item/message, custom_tool_call, event_msg, apply_patch file changes from transcript chunks) - CLI: protocol v2, --agent flag, Codex hook response format - tracevault init --codex installs .codex/hooks.json - AgentBadge component with per-agent icon on session list and detail - Remove old hardcoded extract_file_change/is_file_modifying_tool - 34 adapter tests, UTF-8 safe truncation Co-Authored-By: Claude Opus 4.6 (1M context) --- .gitignore | 1 + README.md | 22 +- crates/tracevault-cli/src/commands/init.rs | 87 ++++ crates/tracevault-cli/src/commands/stream.rs | 20 +- crates/tracevault-cli/src/main.rs | 17 +- .../src/agent_adapter/claude_code.rs | 363 +++++++++++++++ .../src/agent_adapter/codex.rs | 401 +++++++++++++++++ .../src/agent_adapter/default.rs | 46 ++ .../tracevault-core/src/agent_adapter/mod.rs | 89 ++++ crates/tracevault-core/src/lib.rs | 1 + crates/tracevault-core/src/streaming.rs | 44 -- .../tests/agent_adapter_test.rs | 413 ++++++++++++++++++ .../tracevault-core/tests/streaming_test.rs | 62 --- .../src/api/session_detail.rs | 290 +++--------- crates/tracevault-server/src/api/traces_ui.rs | 24 +- crates/tracevault-server/src/lib.rs | 4 + crates/tracevault-server/src/main.rs | 2 + .../tracevault-server/src/service/stream.rs | 85 ++-- web/src/lib/components/AgentBadge.svelte | 69 +++ .../orgs/[slug]/traces/sessions/+page.svelte | 6 +- .../[slug]/traces/sessions/[id]/+page.svelte | 2 + 21 files changed, 1657 insertions(+), 391 deletions(-) create mode 100644 crates/tracevault-core/src/agent_adapter/claude_code.rs create mode 100644 crates/tracevault-core/src/agent_adapter/codex.rs create mode 100644 crates/tracevault-core/src/agent_adapter/default.rs create mode 100644 crates/tracevault-core/src/agent_adapter/mod.rs create mode 100644 crates/tracevault-core/tests/agent_adapter_test.rs create mode 100644 web/src/lib/components/AgentBadge.svelte diff --git a/.gitignore b/.gitignore index 5fafe95..83d255b 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,7 @@ *.swp .DS_Store crates/tracevault-server/data/repos +data/ .playwright-mcp docs/plans/ docs/infra/ diff --git a/README.md b/README.md index 3faf88b..bf45e2d 100644 --- a/README.md +++ b/README.md @@ -6,6 +6,8 @@ AI code governance platform for enterprises. Captures what AI coding agents do in your repos — which files they touch, how many tokens they burn, what tools they call, what percentage of code is AI-generated — then enforces policies and produces tamper-evident audit trails for regulatory compliance. +Supports **Claude Code**, **Codex CLI**, and is extensible to other agents via the AgentAdapter architecture. + Built for financial institutions and regulated industries where AI-generated code needs the same audit rigor as human-written code. [Learn more at VirtusLab](https://virtuslab.com/services/tracevault) @@ -67,7 +69,7 @@ See exactly what AI wrote, line by line. The code browser overlays AI attributio Three Rust crates in a Cargo workspace: - **tracevault-core** — domain types, policy engine (7 condition types), attribution engine (tree-sitter based), secret redactor -- **tracevault-cli** — CLI binary that hooks into Claude Code, captures traces locally, checks policies, pushes to server +- **tracevault-cli** — CLI binary that hooks into Claude Code and Codex CLI, captures traces locally, checks policies, pushes to server - **tracevault-server** — axum HTTP server backed by PostgreSQL with Ed25519 signing, audit logging, RBAC, code browser Plus a SvelteKit web dashboard and a GitHub Action for CI verification. @@ -264,6 +266,19 @@ tracevault init That's it. From this point on, every Claude Code session in this repo is automatically traced — tool calls, file edits, token usage, and model info are captured and streamed to the TraceVault server. When you `git push`, the pre-push hook evaluates policies and uploads traces. +## Using with Codex CLI + +[Codex CLI](https://github.com/openai/codex) (OpenAI's coding agent) is also supported. Initialize with the `--codex` flag to install Codex hooks: + +```sh +npm install -g @openai/codex +cd /path/to/your/repo +tracevault login --server-url https://your-tracevault-server.example.com +tracevault init --codex +``` + +This installs hooks in `.codex/hooks.json` in addition to the Claude Code hooks. Codex sessions are traced including transcript parsing, token usage, and file changes via `apply_patch`. The session detail view shows a Codex badge to distinguish agent types. + ## Keys & Secrets ### Encryption key (`TRACEVAULT_ENCRYPTION_KEY`) @@ -316,10 +331,11 @@ export DATABASE_URL=postgres://user:password@host:5432/tracevault?sslmode=requir | Command | Description | |---------|-------------| -| `tracevault init [--server-url URL]` | Initialize TraceVault in current repo, install pre-push hook and Claude Code hooks | +| `tracevault init [--server-url URL] [--codex]` | Initialize TraceVault in current repo, install hooks (Claude Code by default, `--codex` adds Codex CLI) | | `tracevault login --server-url URL` | Authenticate via device auth flow (opens browser) | | `tracevault logout` | Clear local credentials | -| `tracevault hook --event ` | Handle a Claude Code hook event (reads JSON from stdin) | +| `tracevault hook --event ` | Handle a hook event from any agent (reads JSON from stdin) | +| `tracevault stream --event [--agent ]` | Stream hook events to server (`--agent`: `claude-code` (default), `codex`) | | `tracevault sync` | Sync repo metadata with the server | | `tracevault check` | Evaluate policies against server rules, exit non-zero if blocked | | `tracevault push` | Push collected traces to the server | diff --git a/crates/tracevault-cli/src/commands/init.rs b/crates/tracevault-cli/src/commands/init.rs index 64b5ad0..104229d 100644 --- a/crates/tracevault-cli/src/commands/init.rs +++ b/crates/tracevault-cli/src/commands/init.rs @@ -33,6 +33,7 @@ fn parse_github_org(remote_url: &str) -> Option { pub async fn init_in_directory( project_root: &Path, server_url: Option<&str>, + codex: bool, ) -> Result<(), io::Error> { // Check for git repository if !project_root.join(".git").exists() { @@ -78,6 +79,11 @@ pub async fn init_in_directory( // Install Claude Code hooks into .claude/settings.json install_claude_hooks(project_root)?; + // Install Codex hooks into .codex/hooks.json + if codex { + install_codex_hooks(project_root)?; + } + // Install git hooks install_git_hook(project_root)?; install_post_commit_hook(project_root)?; @@ -333,6 +339,87 @@ pub fn tracevault_hooks() -> serde_json::Value { }) } +fn codex_hooks() -> serde_json::Value { + serde_json::json!({ + "hooks": { + "SessionStart": [{ + "matcher": "startup|resume", + "hooks": [{ + "type": "command", + "command": "tracevault stream --agent codex --event session-start", + "timeout": 10, + "statusMessage": "TraceVault: streaming session start" + }] + }], + "PreToolUse": [{ + "matcher": "Bash", + "hooks": [{ + "type": "command", + "command": "tracevault stream --agent codex --event pre-tool-use", + "timeout": 10, + "statusMessage": "TraceVault: streaming pre-tool event" + }] + }], + "PostToolUse": [{ + "matcher": "Bash", + "hooks": [{ + "type": "command", + "command": "tracevault stream --agent codex --event post-tool-use", + "timeout": 10, + "statusMessage": "TraceVault: streaming post-tool event" + }] + }], + "Stop": [{ + "hooks": [{ + "type": "command", + "command": "tracevault stream --agent codex --event stop", + "timeout": 10, + "statusMessage": "TraceVault: finalizing session" + }] + }] + } + }) +} + +fn install_codex_hooks(project_root: &Path) -> Result<(), io::Error> { + let codex_dir = project_root.join(".codex"); + fs::create_dir_all(&codex_dir)?; + + let hooks_path = codex_dir.join("hooks.json"); + let mut config: serde_json::Value = if hooks_path.exists() { + let content = fs::read_to_string(&hooks_path)?; + serde_json::from_str(&content).map_err(|e| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("Failed to parse .codex/hooks.json: {e}"), + ) + })? + } else { + serde_json::json!({}) + }; + + let hooks = codex_hooks(); + + // Merge hooks into existing config + let config_obj = config.as_object_mut().ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidData, + ".codex/hooks.json is not a JSON object", + ) + })?; + + // Set hooks key from our template + if let Some(hooks_value) = hooks.get("hooks") { + config_obj.insert("hooks".to_string(), hooks_value.clone()); + } + + let formatted = serde_json::to_string_pretty(&config) + .map_err(|e| io::Error::other(format!("Failed to serialize hooks: {e}")))?; + fs::write(&hooks_path, formatted)?; + + Ok(()) +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/tracevault-cli/src/commands/stream.rs b/crates/tracevault-cli/src/commands/stream.rs index c372c6c..22d4720 100644 --- a/crates/tracevault-cli/src/commands/stream.rs +++ b/crates/tracevault-cli/src/commands/stream.rs @@ -85,6 +85,7 @@ pub fn drain_pending(pending_path: &Path) -> Result, io::Error> { pub async fn run_stream( project_root: &Path, event_type: &str, + agent: &str, ) -> Result<(), Box> { // 1. Read HookEvent from stdin let mut input = String::new(); @@ -108,15 +109,16 @@ pub async fn run_stream( let (transcript_lines, new_offset) = read_new_transcript_lines(transcript_path, &offset_path)?; // 5. Build StreamEventRequest + // Claude Code sends "notification" for SessionStart, Codex sends "session-start" let stream_event_type = match event_type { - "notification" => StreamEventType::SessionStart, + "notification" | "session-start" => StreamEventType::SessionStart, "stop" => StreamEventType::SessionEnd, _ => StreamEventType::ToolUse, }; let req = StreamEventRequest { - protocol_version: 1, - tool: Some("claude-code".to_string()), + protocol_version: 2, + tool: Some(agent.to_string()), event_type: stream_event_type, session_id: hook_event.session_id.clone(), timestamp: chrono::Utc::now(), @@ -193,9 +195,15 @@ pub async fn run_stream( } } - // 12. Always print HookResponse::allow() to stdout - let response = HookResponse::allow(); - println!("{}", serde_json::to_string(&response)?); + // 12. Always print hook response to stdout + // Codex expects empty JSON {}, Claude Code expects {"suppress_output": true} + match agent { + "codex" => println!("{{}}"), + _ => { + let response = HookResponse::allow(); + println!("{}", serde_json::to_string(&response)?); + } + } Ok(()) } diff --git a/crates/tracevault-cli/src/main.rs b/crates/tracevault-cli/src/main.rs index 2ce3528..bf18460 100644 --- a/crates/tracevault-cli/src/main.rs +++ b/crates/tracevault-cli/src/main.rs @@ -15,6 +15,9 @@ enum Cli { /// TraceVault server URL for repo registration #[arg(long)] server_url: Option, + /// Also install Codex CLI hooks (.codex/hooks.json) + #[arg(long)] + codex: bool, }, /// Show current session status Status, @@ -27,6 +30,9 @@ enum Cli { Stream { #[arg(long)] event: String, + /// AI coding agent name (claude-code, codex) + #[arg(long, default_value = "claude-code")] + agent: String, }, /// Check session policies before pushing Check, @@ -63,12 +69,15 @@ enum Cli { async fn main() { let cli = Cli::parse(); match cli { - Cli::Init { server_url } => { + Cli::Init { server_url, codex } => { let cwd = env::current_dir().expect("Cannot determine current directory"); - match commands::init::init_in_directory(&cwd, server_url.as_deref()).await { + match commands::init::init_in_directory(&cwd, server_url.as_deref(), codex).await { Ok(()) => { println!("TraceVault initialized in {}", cwd.display()); println!("Claude Code hooks installed in .claude/settings.json"); + if codex { + println!("Codex hooks installed in .codex/hooks.json"); + } println!("Git pre-push hook installed"); } Err(e) => eprintln!("Error: {e}"), @@ -81,9 +90,9 @@ async fn main() { eprintln!("Hook error: {e}"); } } - Cli::Stream { event } => { + Cli::Stream { event, agent } => { let cwd = env::current_dir().expect("Cannot determine current directory"); - if let Err(e) = commands::stream::run_stream(&cwd, &event).await { + if let Err(e) = commands::stream::run_stream(&cwd, &event, &agent).await { eprintln!("Stream error: {e}"); } } diff --git a/crates/tracevault-core/src/agent_adapter/claude_code.rs b/crates/tracevault-core/src/agent_adapter/claude_code.rs new file mode 100644 index 0000000..b412b03 --- /dev/null +++ b/crates/tracevault-core/src/agent_adapter/claude_code.rs @@ -0,0 +1,363 @@ +use sha2::{Digest, Sha256}; + +use crate::streaming::{ExtractedFileChange, StreamEventType}; + +use super::{AgentAdapter, ParsedTranscriptRecord, TokenUsage}; + +pub struct ClaudeCodeAdapter; + +impl AgentAdapter for ClaudeCodeAdapter { + fn name(&self) -> &str { + "claude-code" + } + + fn map_event_type(&self, hook_event_name: &str) -> StreamEventType { + match hook_event_name { + "SessionStart" => StreamEventType::SessionStart, + "Stop" => StreamEventType::SessionEnd, + _ => StreamEventType::ToolUse, + } + } + + fn is_file_modifying(&self, tool_name: &str) -> bool { + matches!(tool_name, "Write" | "Edit" | "Bash") + } + + fn extract_file_changes( + &self, + tool_name: &str, + tool_input: &serde_json::Value, + ) -> Vec { + match tool_name { + "Write" => { + let file_path = match tool_input.get("file_path").and_then(|v| v.as_str()) { + Some(p) => p.to_string(), + None => return Vec::new(), + }; + let content = match tool_input.get("content").and_then(|v| v.as_str()) { + Some(c) => c, + None => return Vec::new(), + }; + let mut hasher = Sha256::new(); + hasher.update(content.as_bytes()); + let hash = format!("{:x}", hasher.finalize()); + let diff_text = content + .lines() + .map(|line| format!("+{}", line)) + .collect::>() + .join("\n"); + vec![ExtractedFileChange { + file_path, + change_type: "create".to_string(), + diff_text: Some(diff_text), + content_hash: Some(hash), + }] + } + "Edit" => { + let file_path = match tool_input.get("file_path").and_then(|v| v.as_str()) { + Some(p) => p.to_string(), + None => return Vec::new(), + }; + let old_string = match tool_input.get("old_string").and_then(|v| v.as_str()) { + Some(s) => s, + None => return Vec::new(), + }; + let new_string = match tool_input.get("new_string").and_then(|v| v.as_str()) { + Some(s) => s, + None => return Vec::new(), + }; + let diff_text = format!("--- {}\n+++ {}", old_string, new_string); + vec![ExtractedFileChange { + file_path, + change_type: "edit".to_string(), + diff_text: Some(diff_text), + content_hash: None, + }] + } + _ => Vec::new(), + } + } + + fn extract_token_usage(&self, chunk: &serde_json::Value) -> Option { + let usage = chunk.get("message")?.get("usage")?; + Some(TokenUsage { + input_tokens: usage.get("input_tokens")?.as_i64()?, + output_tokens: usage.get("output_tokens")?.as_i64()?, + cache_read_tokens: usage + .get("cache_read_input_tokens") + .and_then(|v| v.as_i64()) + .unwrap_or(0), + cache_write_tokens: usage + .get("cache_creation_input_tokens") + .and_then(|v| v.as_i64()) + .unwrap_or(0), + }) + } + + fn extract_model(&self, chunk: &serde_json::Value) -> Option { + chunk + .get("message")? + .get("model")? + .as_str() + .map(|s| s.to_string()) + } + + fn parse_transcript_record(&self, chunk: &serde_json::Value) -> Option { + let record_type = chunk.get("type")?.as_str()?.to_string(); + let timestamp = chunk + .get("timestamp") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + match record_type.as_str() { + "assistant" => self.parse_assistant_record(chunk, record_type, timestamp), + "user" => self.parse_user_record(chunk, record_type, timestamp), + "progress" => self.parse_progress_record(chunk, record_type, timestamp), + "system" => self.parse_system_record(chunk, record_type, timestamp), + _ => Some(ParsedTranscriptRecord { + record_type, + timestamp, + content_types: Vec::new(), + tool_name: None, + text: None, + raw_input_tokens: None, + raw_output_tokens: None, + raw_cache_read_tokens: None, + raw_cache_write_tokens: None, + model: None, + }), + } + } +} + +impl ClaudeCodeAdapter { + fn parse_assistant_record( + &self, + chunk: &serde_json::Value, + record_type: String, + timestamp: Option, + ) -> Option { + let message = chunk.get("message")?; + let model = message + .get("model") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + let mut content_types = Vec::new(); + let mut text_parts = Vec::new(); + let mut first_tool_name: Option = None; + + if let Some(content) = message.get("content").and_then(|v| v.as_array()) { + for block in content { + if let Some(block_type) = block.get("type").and_then(|v| v.as_str()) { + if !content_types.contains(&block_type.to_string()) { + content_types.push(block_type.to_string()); + } + match block_type { + "text" => { + if let Some(t) = block.get("text").and_then(|v| v.as_str()) { + text_parts.push(t.to_string()); + } + } + "thinking" => { + if let Some(t) = block.get("thinking").and_then(|v| v.as_str()) { + text_parts.push(t.to_string()); + } + } + "tool_use" => { + if first_tool_name.is_none() { + first_tool_name = block + .get("name") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + } + } + _ => {} + } + } + } + } + + let usage = message.get("usage"); + let raw_input_tokens = usage + .and_then(|u| u.get("input_tokens")) + .and_then(|v| v.as_i64()); + let raw_output_tokens = usage + .and_then(|u| u.get("output_tokens")) + .and_then(|v| v.as_i64()); + let raw_cache_read_tokens = usage + .and_then(|u| u.get("cache_read_input_tokens")) + .and_then(|v| v.as_i64()); + let raw_cache_write_tokens = usage + .and_then(|u| u.get("cache_creation_input_tokens")) + .and_then(|v| v.as_i64()); + + let text = if text_parts.is_empty() { + None + } else { + Some(text_parts.join("\n")) + }; + + Some(ParsedTranscriptRecord { + record_type, + timestamp, + content_types, + tool_name: first_tool_name, + text, + raw_input_tokens, + raw_output_tokens, + raw_cache_read_tokens, + raw_cache_write_tokens, + model, + }) + } + + fn parse_user_record( + &self, + chunk: &serde_json::Value, + record_type: String, + timestamp: Option, + ) -> Option { + let mut content_types = Vec::new(); + let mut text_parts = Vec::new(); + let mut tool_name: Option = None; + + // Check for toolUseResult (e.g. Read, Glob, Bash results) + if let Some(tool_result) = chunk.get("toolUseResult") { + if let Some(file_info) = tool_result.get("file") { + let file_path = file_info + .get("filePath") + .and_then(|v| v.as_str()) + .unwrap_or("unknown"); + tool_name = Some(format!("Read: {}", file_path)); + } else if let Some(glob_info) = tool_result.get("glob") { + let pattern = glob_info + .get("pattern") + .and_then(|v| v.as_str()) + .unwrap_or("unknown"); + tool_name = Some(format!("Glob: {}", pattern)); + } else if let Some(bash_info) = tool_result.get("bash") { + let command = bash_info + .get("command") + .and_then(|v| v.as_str()) + .unwrap_or("unknown"); + tool_name = Some(format!("Bash: {}", command)); + } + } + + // Handle message.content as either a string or an array + if let Some(message) = chunk.get("message") { + if let Some(content) = message.get("content") { + if let Some(text) = content.as_str() { + text_parts.push(text.to_string()); + content_types.push("text".to_string()); + } else if let Some(arr) = content.as_array() { + for block in arr { + if let Some(block_type) = block.get("type").and_then(|v| v.as_str()) { + if !content_types.contains(&block_type.to_string()) { + content_types.push(block_type.to_string()); + } + match block_type { + "tool_result" | "text" => { + if let Some(t) = block.get("text").and_then(|v| v.as_str()) { + text_parts.push(t.to_string()); + } + } + _ => {} + } + } + } + } + } + } + + let text = if text_parts.is_empty() { + None + } else { + Some(text_parts.join("\n")) + }; + + Some(ParsedTranscriptRecord { + record_type, + timestamp, + content_types, + tool_name, + text, + raw_input_tokens: None, + raw_output_tokens: None, + raw_cache_read_tokens: None, + raw_cache_write_tokens: None, + model: None, + }) + } + + fn parse_progress_record( + &self, + chunk: &serde_json::Value, + record_type: String, + timestamp: Option, + ) -> Option { + let data = chunk.get("data"); + let hook_name = data + .and_then(|d| d.get("hookName")) + .and_then(|v| v.as_str()) + .unwrap_or("unknown"); + let hook_event = data + .and_then(|d| d.get("hookEvent")) + .and_then(|v| v.as_str()) + .unwrap_or("unknown"); + + let text = format!("{}: {}", hook_event, hook_name); + + Some(ParsedTranscriptRecord { + record_type, + timestamp, + content_types: Vec::new(), + tool_name: None, + text: Some(text), + raw_input_tokens: None, + raw_output_tokens: None, + raw_cache_read_tokens: None, + raw_cache_write_tokens: None, + model: None, + }) + } + + fn parse_system_record( + &self, + chunk: &serde_json::Value, + record_type: String, + timestamp: Option, + ) -> Option { + let subtype = chunk.get("subtype").and_then(|v| v.as_str()); + + let text = match subtype { + Some("turn_duration") => { + let duration_ms = chunk + .get("durationMs") + .and_then(|v| v.as_f64()) + .unwrap_or(0.0); + let seconds = duration_ms / 1000.0; + Some(format!("Turn duration: {:.1}s", seconds)) + } + Some("stop_hook_summary") => { + let hook_count = chunk.get("hookCount").and_then(|v| v.as_i64()).unwrap_or(0); + Some(format!("Stop hooks executed: {}", hook_count)) + } + _ => None, + }; + + Some(ParsedTranscriptRecord { + record_type, + timestamp, + content_types: Vec::new(), + tool_name: None, + text, + raw_input_tokens: None, + raw_output_tokens: None, + raw_cache_read_tokens: None, + raw_cache_write_tokens: None, + model: None, + }) + } +} diff --git a/crates/tracevault-core/src/agent_adapter/codex.rs b/crates/tracevault-core/src/agent_adapter/codex.rs new file mode 100644 index 0000000..7601bc0 --- /dev/null +++ b/crates/tracevault-core/src/agent_adapter/codex.rs @@ -0,0 +1,401 @@ +use sha2::{Digest, Sha256}; + +use crate::streaming::{ExtractedFileChange, StreamEventType}; + +use super::{AgentAdapter, ParsedTranscriptRecord, TokenUsage}; + +/// Adapter for OpenAI Codex CLI. +/// +/// Codex file modifications come exclusively through transcript chunks +/// (custom_tool_call with apply_patch), NOT through hook ToolUse events. +/// The hook events only carry shell commands like `pwd`, `git status`, etc. +pub struct CodexAdapter; + +impl AgentAdapter for CodexAdapter { + fn name(&self) -> &str { + "codex" + } + + fn map_event_type(&self, hook_event_name: &str) -> StreamEventType { + match hook_event_name { + "SessionStart" => StreamEventType::SessionStart, + "Stop" => StreamEventType::SessionEnd, + _ => StreamEventType::ToolUse, + } + } + + /// Codex hook events never carry file-modifying tool calls. + /// File changes are extracted from transcript via `extract_file_changes_from_transcript`. + fn is_file_modifying(&self, _tool_name: &str) -> bool { + false + } + + /// Not used for Codex — file changes come from transcript, not hook events. + fn extract_file_changes( + &self, + _tool_name: &str, + _tool_input: &serde_json::Value, + ) -> Vec { + vec![] + } + + /// Extract file changes from Codex transcript chunks. + /// Handles `response_item` with `payload.type: "custom_tool_call"` and `name: "apply_patch"`. + fn extract_file_changes_from_transcript( + &self, + chunk: &serde_json::Value, + ) -> Vec { + let payload = match chunk.get("payload") { + Some(p) => p, + None => return vec![], + }; + + let payload_type = payload.get("type").and_then(|v| v.as_str()).unwrap_or(""); + if payload_type != "custom_tool_call" { + return vec![]; + } + + let name = payload.get("name").and_then(|v| v.as_str()).unwrap_or(""); + if name != "apply_patch" { + return vec![]; + } + + let input = match payload.get("input").and_then(|v| v.as_str()) { + Some(s) => s, + None => return vec![], + }; + + parse_codex_patch(input) + } + + fn extract_token_usage(&self, chunk: &serde_json::Value) -> Option { + let top_type = chunk.get("type")?.as_str()?; + if top_type != "event_msg" { + return None; + } + let payload = chunk.get("payload")?; + let payload_type = payload.get("type")?.as_str()?; + if payload_type != "token_count" { + return None; + } + let usage = payload.get("info")?.get("last_token_usage")?; + Some(TokenUsage { + input_tokens: usage.get("input_tokens")?.as_i64()?, + output_tokens: usage.get("output_tokens")?.as_i64()?, + cache_read_tokens: usage + .get("cached_input_tokens") + .and_then(|v| v.as_i64()) + .unwrap_or(0), + cache_write_tokens: 0, + }) + } + + fn extract_model(&self, chunk: &serde_json::Value) -> Option { + let top_type = chunk.get("type")?.as_str()?; + if top_type != "turn_context" { + return None; + } + chunk + .get("payload")? + .get("model")? + .as_str() + .map(|s| s.to_string()) + } + + fn parse_transcript_record(&self, chunk: &serde_json::Value) -> Option { + let top_type = chunk.get("type")?.as_str()?; + let timestamp = chunk + .get("timestamp") + .and_then(|v| v.as_str()) + .map(String::from); + + match top_type { + "event_msg" => self.parse_event_msg(chunk, ×tamp), + "response_item" => self.parse_response_item(chunk, ×tamp), + // turn_context, session_meta — ingestion-only, not for display + _ => None, + } + } +} + +impl CodexAdapter { + fn parse_event_msg( + &self, + chunk: &serde_json::Value, + timestamp: &Option, + ) -> Option { + let payload = chunk.get("payload")?; + let payload_type = payload.get("type")?.as_str()?; + + match payload_type { + "agent_message" => { + let content = payload + .get("content") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + Some(ParsedTranscriptRecord { + record_type: "assistant".to_string(), + timestamp: timestamp.clone(), + content_types: vec!["text".to_string()], + tool_name: None, + text: content, + raw_input_tokens: None, + raw_output_tokens: None, + raw_cache_read_tokens: None, + raw_cache_write_tokens: None, + model: None, + }) + } + "user_message" => { + let content = payload + .get("content") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + Some(ParsedTranscriptRecord { + record_type: "user".to_string(), + timestamp: timestamp.clone(), + content_types: vec!["text".to_string()], + tool_name: None, + text: content, + raw_input_tokens: None, + raw_output_tokens: None, + raw_cache_read_tokens: None, + raw_cache_write_tokens: None, + model: None, + }) + } + // token_count, task_started — ingestion-only + _ => None, + } + } + + fn parse_response_item( + &self, + chunk: &serde_json::Value, + timestamp: &Option, + ) -> Option { + let payload = chunk.get("payload")?; + let payload_type = payload.get("type")?.as_str()?; + + match payload_type { + "local_shell_call" => { + let command = payload + .get("command") + .and_then(|v| v.as_str()) + .unwrap_or(""); + let output = payload.get("output").and_then(|v| v.as_str()).unwrap_or(""); + let text = format!("$ {}\n{}", command, output); + Some(ParsedTranscriptRecord { + record_type: "assistant".to_string(), + timestamp: timestamp.clone(), + content_types: vec!["tool_use".to_string()], + tool_name: Some("Bash".to_string()), + text: Some(text), + raw_input_tokens: None, + raw_output_tokens: None, + raw_cache_read_tokens: None, + raw_cache_write_tokens: None, + model: None, + }) + } + "message" => { + let role = payload.get("role")?.as_str()?; + // Skip system/developer messages (permissions, instructions) + if role == "developer" { + return None; + } + let record_type = if role == "assistant" { + "assistant" + } else { + "user" + }; + let text = payload + .get("content") + .and_then(|v| v.as_array()) + .map(|arr| { + arr.iter() + .filter_map(|block| { + let block_type = block.get("type").and_then(|v| v.as_str())?; + if block_type == "input_text" || block_type == "output_text" { + let t = block.get("text").and_then(|v| v.as_str())?; + // Skip system prompts (XML tags in user messages) + if t.starts_with('<') && role == "user" { + return None; + } + Some(t.to_string()) + } else { + None + } + }) + .collect::>() + .join("\n\n") + }) + .filter(|s| !s.is_empty()); + // Skip if no meaningful text + text.as_ref()?; + Some(ParsedTranscriptRecord { + record_type: record_type.to_string(), + timestamp: timestamp.clone(), + content_types: vec!["text".to_string()], + tool_name: None, + text, + raw_input_tokens: None, + raw_output_tokens: None, + raw_cache_read_tokens: None, + raw_cache_write_tokens: None, + model: None, + }) + } + "custom_tool_call" => { + let name = payload + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("tool"); + let input = payload.get("input").and_then(|v| v.as_str()).unwrap_or(""); + // Truncate long patches for display (char-safe to avoid UTF-8 panic) + let display_input = if input.len() > 500 { + let truncated: String = input.chars().take(500).collect(); + format!("{}...", truncated) + } else { + input.to_string() + }; + Some(ParsedTranscriptRecord { + record_type: "assistant".to_string(), + timestamp: timestamp.clone(), + content_types: vec!["tool_use".to_string()], + tool_name: Some(name.to_string()), + text: Some(display_input), + raw_input_tokens: None, + raw_output_tokens: None, + raw_cache_read_tokens: None, + raw_cache_write_tokens: None, + model: None, + }) + } + // reasoning — encrypted, skip + _ => None, + } + } +} + +/// Parse Codex's custom apply_patch format into file changes. +pub fn parse_codex_patch(patch: &str) -> Vec { + let mut changes = Vec::new(); + let mut current_file: Option = None; + let mut current_type: Option = None; + let mut current_lines: Vec = Vec::new(); + + for line in patch.lines() { + if line == "*** Begin Patch" || line == "*** End Patch" { + flush_pending( + &mut changes, + &mut current_file, + &mut current_type, + &mut current_lines, + ); + continue; + } + + if let Some(path) = line.strip_prefix("*** Add File: ") { + flush_pending( + &mut changes, + &mut current_file, + &mut current_type, + &mut current_lines, + ); + current_file = Some(path.to_string()); + current_type = Some("create".to_string()); + } else if let Some(path) = line.strip_prefix("*** Update File: ") { + flush_pending( + &mut changes, + &mut current_file, + &mut current_type, + &mut current_lines, + ); + current_file = Some(path.to_string()); + current_type = Some("edit".to_string()); + } else if let Some(path) = line.strip_prefix("*** Delete File: ") { + flush_pending( + &mut changes, + &mut current_file, + &mut current_type, + &mut current_lines, + ); + current_file = Some(path.to_string()); + current_type = Some("delete".to_string()); + } else if current_file.is_some() { + current_lines.push(line.to_string()); + } + } + + flush_pending( + &mut changes, + &mut current_file, + &mut current_type, + &mut current_lines, + ); + changes +} + +fn flush_pending( + changes: &mut Vec, + file: &mut Option, + kind: &mut Option, + lines: &mut Vec, +) { + if let (Some(file_path), Some(change_type)) = (file.take(), kind.take()) { + changes.push(build_file_change(&file_path, &change_type, lines)); + lines.clear(); + } +} + +fn build_file_change(file_path: &str, change_type: &str, lines: &[String]) -> ExtractedFileChange { + match change_type { + "create" => { + let content: String = lines + .iter() + .map(|l| l.strip_prefix('+').unwrap_or(l)) + .collect::>() + .join("\n"); + let mut hasher = Sha256::new(); + hasher.update(content.as_bytes()); + let hash = format!("{:x}", hasher.finalize()); + let diff_text = lines.join("\n"); + ExtractedFileChange { + file_path: file_path.to_string(), + change_type: "create".to_string(), + diff_text: if diff_text.is_empty() { + None + } else { + Some(diff_text) + }, + content_hash: Some(hash), + } + } + "edit" => { + let diff_text = lines.join("\n"); + ExtractedFileChange { + file_path: file_path.to_string(), + change_type: "edit".to_string(), + diff_text: if diff_text.is_empty() { + None + } else { + Some(diff_text) + }, + content_hash: None, + } + } + "delete" => ExtractedFileChange { + file_path: file_path.to_string(), + change_type: "delete".to_string(), + diff_text: None, + content_hash: None, + }, + _ => ExtractedFileChange { + file_path: file_path.to_string(), + change_type: change_type.to_string(), + diff_text: None, + content_hash: None, + }, + } +} diff --git a/crates/tracevault-core/src/agent_adapter/default.rs b/crates/tracevault-core/src/agent_adapter/default.rs new file mode 100644 index 0000000..c4e3580 --- /dev/null +++ b/crates/tracevault-core/src/agent_adapter/default.rs @@ -0,0 +1,46 @@ +use crate::streaming::{ExtractedFileChange, StreamEventType}; + +use super::{AgentAdapter, ParsedTranscriptRecord, TokenUsage}; + +pub struct DefaultAdapter; + +impl AgentAdapter for DefaultAdapter { + fn name(&self) -> &str { + "default" + } + + fn map_event_type(&self, hook_event_name: &str) -> StreamEventType { + match hook_event_name { + "SessionStart" => StreamEventType::SessionStart, + "Stop" | "SessionEnd" => StreamEventType::SessionEnd, + _ => StreamEventType::ToolUse, + } + } + + fn is_file_modifying(&self, _tool_name: &str) -> bool { + false + } + + fn extract_file_changes( + &self, + _tool_name: &str, + _tool_input: &serde_json::Value, + ) -> Vec { + Vec::new() + } + + fn extract_token_usage(&self, _chunk: &serde_json::Value) -> Option { + None + } + + fn extract_model(&self, _chunk: &serde_json::Value) -> Option { + None + } + + fn parse_transcript_record( + &self, + _chunk: &serde_json::Value, + ) -> Option { + None + } +} diff --git a/crates/tracevault-core/src/agent_adapter/mod.rs b/crates/tracevault-core/src/agent_adapter/mod.rs new file mode 100644 index 0000000..d8660f0 --- /dev/null +++ b/crates/tracevault-core/src/agent_adapter/mod.rs @@ -0,0 +1,89 @@ +pub mod claude_code; +pub mod codex; +mod default; + +use serde::Serialize; +use std::collections::HashMap; +use std::sync::Arc; + +use crate::streaming::{ExtractedFileChange, StreamEventType}; + +use self::default::DefaultAdapter; + +#[derive(Debug, Clone, Default)] +pub struct TokenUsage { + pub input_tokens: i64, + pub output_tokens: i64, + pub cache_read_tokens: i64, + pub cache_write_tokens: i64, +} + +#[derive(Debug, Clone, Serialize)] +pub struct ParsedTranscriptRecord { + pub record_type: String, + pub timestamp: Option, + pub content_types: Vec, + pub tool_name: Option, + pub text: Option, + pub raw_input_tokens: Option, + pub raw_output_tokens: Option, + pub raw_cache_read_tokens: Option, + pub raw_cache_write_tokens: Option, + pub model: Option, +} + +pub trait AgentAdapter: Send + Sync { + fn name(&self) -> &str; + fn map_event_type(&self, hook_event_name: &str) -> StreamEventType; + fn is_file_modifying(&self, tool_name: &str) -> bool; + /// Extract file changes from a hook tool event (tool_name + tool_input) + fn extract_file_changes( + &self, + tool_name: &str, + tool_input: &serde_json::Value, + ) -> Vec; + /// Extract file changes from a transcript chunk (e.g. Codex custom_tool_call with apply_patch). + /// Default: no extraction. Override for agents whose file ops appear in transcript, not hook events. + fn extract_file_changes_from_transcript( + &self, + _chunk: &serde_json::Value, + ) -> Vec { + vec![] + } + fn extract_token_usage(&self, chunk: &serde_json::Value) -> Option; + fn extract_model(&self, chunk: &serde_json::Value) -> Option; + fn parse_transcript_record(&self, chunk: &serde_json::Value) -> Option; +} + +pub struct AgentAdapterRegistry { + adapters: HashMap>, + default: Arc, +} + +impl AgentAdapterRegistry { + pub fn new() -> Self { + let mut adapters: HashMap> = HashMap::new(); + adapters.insert( + "claude-code".to_string(), + Arc::new(claude_code::ClaudeCodeAdapter), + ); + adapters.insert("codex".to_string(), Arc::new(codex::CodexAdapter)); + Self { + adapters, + default: Arc::new(DefaultAdapter), + } + } + + pub fn get(&self, name: &str) -> &dyn AgentAdapter { + self.adapters + .get(name) + .map(|a| a.as_ref()) + .unwrap_or(self.default.as_ref()) + } +} + +impl Default for AgentAdapterRegistry { + fn default() -> Self { + Self::new() + } +} diff --git a/crates/tracevault-core/src/lib.rs b/crates/tracevault-core/src/lib.rs index d1129d6..73101ae 100644 --- a/crates/tracevault-core/src/lib.rs +++ b/crates/tracevault-core/src/lib.rs @@ -1,3 +1,4 @@ +pub mod agent_adapter; pub mod code_nav; pub mod diff; pub mod extensions; diff --git a/crates/tracevault-core/src/streaming.rs b/crates/tracevault-core/src/streaming.rs index 62819fe..d366398 100644 --- a/crates/tracevault-core/src/streaming.rs +++ b/crates/tracevault-core/src/streaming.rs @@ -1,6 +1,5 @@ use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; -use sha2::{Digest, Sha256}; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] #[serde(rename_all = "snake_case")] @@ -74,46 +73,3 @@ pub struct ExtractedFileChange { pub diff_text: Option, pub content_hash: Option, } - -pub fn is_file_modifying_tool(tool_name: &str) -> bool { - matches!(tool_name, "Write" | "Edit" | "Bash") -} - -pub fn extract_file_change( - tool_name: &str, - tool_input: &serde_json::Value, -) -> Option { - match tool_name { - "Write" => { - let file_path = tool_input.get("file_path")?.as_str()?.to_string(); - let content = tool_input.get("content")?.as_str()?; - let mut hasher = Sha256::new(); - hasher.update(content.as_bytes()); - let hash = format!("{:x}", hasher.finalize()); - let diff = content - .lines() - .map(|l| format!("+{l}")) - .collect::>() - .join("\n"); - Some(ExtractedFileChange { - file_path, - change_type: "create".to_string(), - diff_text: Some(diff), - content_hash: Some(hash), - }) - } - "Edit" => { - let file_path = tool_input.get("file_path")?.as_str()?.to_string(); - let old_string = tool_input.get("old_string")?.as_str()?; - let new_string = tool_input.get("new_string")?.as_str()?; - let diff = format!("--- {old_string}\n+++ {new_string}"); - Some(ExtractedFileChange { - file_path, - change_type: "edit".to_string(), - diff_text: Some(diff), - content_hash: None, - }) - } - _ => None, - } -} diff --git a/crates/tracevault-core/tests/agent_adapter_test.rs b/crates/tracevault-core/tests/agent_adapter_test.rs new file mode 100644 index 0000000..2cc99dd --- /dev/null +++ b/crates/tracevault-core/tests/agent_adapter_test.rs @@ -0,0 +1,413 @@ +use serde_json::json; +use tracevault_core::agent_adapter::{AgentAdapterRegistry, TokenUsage}; +use tracevault_core::streaming::StreamEventType; + +#[test] +fn registry_unknown_agent_returns_default() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("unknown-agent"); + assert_eq!(adapter.name(), "default"); +} + +#[test] +fn default_adapter_extract_token_usage_returns_none() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("nope"); + let chunk = + serde_json::json!({"type": "assistant", "message": {"usage": {"input_tokens": 100}}}); + assert!(adapter.extract_token_usage(&chunk).is_none()); +} + +#[test] +fn registry_dispatches_to_claude_code() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("claude-code"); + assert_eq!(adapter.name(), "claude-code"); +} + +#[test] +fn claude_code_map_event_types() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("claude-code"); + assert!(matches!( + adapter.map_event_type("SessionStart"), + StreamEventType::SessionStart + )); + assert!(matches!( + adapter.map_event_type("Stop"), + StreamEventType::SessionEnd + )); + assert!(matches!( + adapter.map_event_type("PostToolUse"), + StreamEventType::ToolUse + )); +} + +#[test] +fn claude_code_extract_file_change_write() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("claude-code"); + let input = json!({"file_path": "src/main.rs", "content": "fn main() {}"}); + let changes = adapter.extract_file_changes("Write", &input); + assert_eq!(changes.len(), 1); + assert_eq!(changes[0].file_path, "src/main.rs"); + assert_eq!(changes[0].change_type, "create"); + assert!(changes[0].content_hash.is_some()); +} + +#[test] +fn claude_code_extract_file_change_edit() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("claude-code"); + let input = json!({"file_path": "src/lib.rs", "old_string": "old", "new_string": "new"}); + let changes = adapter.extract_file_changes("Edit", &input); + assert_eq!(changes.len(), 1); + assert_eq!(changes[0].change_type, "edit"); +} + +#[test] +fn claude_code_read_returns_empty() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("claude-code"); + let input = json!({"file_path": "src/lib.rs"}); + assert!(adapter.extract_file_changes("Read", &input).is_empty()); +} + +#[test] +fn claude_code_is_file_modifying() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("claude-code"); + assert!(adapter.is_file_modifying("Write")); + assert!(adapter.is_file_modifying("Edit")); + assert!(adapter.is_file_modifying("Bash")); + assert!(!adapter.is_file_modifying("Read")); +} + +#[test] +fn claude_code_extract_token_usage() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("claude-code"); + let chunk = json!({ + "type": "assistant", + "message": { + "usage": { + "input_tokens": 1000, + "output_tokens": 200, + "cache_read_input_tokens": 500, + "cache_creation_input_tokens": 100 + } + } + }); + let usage = adapter.extract_token_usage(&chunk).unwrap(); + assert_eq!(usage.input_tokens, 1000); + assert_eq!(usage.output_tokens, 200); + assert_eq!(usage.cache_read_tokens, 500); + assert_eq!(usage.cache_write_tokens, 100); +} + +#[test] +fn claude_code_extract_model() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("claude-code"); + let chunk = json!({"type": "assistant", "message": {"model": "claude-opus-4-6"}}); + assert_eq!( + adapter.extract_model(&chunk).as_deref(), + Some("claude-opus-4-6") + ); +} + +#[test] +fn claude_code_parse_assistant_record() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("claude-code"); + let chunk = json!({ + "type": "assistant", + "timestamp": "2026-03-23T13:17:16Z", + "message": { + "model": "claude-opus-4-6", + "content": [ + {"type": "text", "text": "Hello world"}, + {"type": "tool_use", "name": "Write", "input": {}} + ], + "usage": { + "input_tokens": 100, "output_tokens": 50, + "cache_read_input_tokens": 0, "cache_creation_input_tokens": 0 + } + } + }); + let record = adapter.parse_transcript_record(&chunk).unwrap(); + assert_eq!(record.record_type, "assistant"); + assert_eq!(record.model.as_deref(), Some("claude-opus-4-6")); + assert!(record.text.as_ref().unwrap().contains("Hello world")); + assert!(record.content_types.contains(&"text".to_string())); + assert!(record.content_types.contains(&"tool_use".to_string())); + assert_eq!(record.tool_name.as_deref(), Some("Write")); + assert_eq!(record.raw_input_tokens, Some(100)); +} + +#[test] +fn claude_code_parse_user_record() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("claude-code"); + let chunk = json!({"type": "user", "timestamp": "2026-03-23T13:17:00Z", "message": {"content": "Fix the bug"}}); + let record = adapter.parse_transcript_record(&chunk).unwrap(); + assert_eq!(record.record_type, "user"); + assert_eq!(record.text.as_deref(), Some("Fix the bug")); +} + +#[test] +fn claude_code_parse_user_tool_result() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("claude-code"); + let chunk = json!({"type": "user", "toolUseResult": {"file": {"filePath": "src/main.rs"}}}); + let record = adapter.parse_transcript_record(&chunk).unwrap(); + assert_eq!(record.tool_name.as_deref(), Some("Read: src/main.rs")); +} + +#[test] +fn claude_code_parse_progress_record() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("claude-code"); + let chunk = + json!({"type": "progress", "data": {"hookName": "tracevault", "hookEvent": "PostToolUse"}}); + let record = adapter.parse_transcript_record(&chunk).unwrap(); + assert_eq!(record.record_type, "progress"); + assert_eq!(record.text.as_deref(), Some("PostToolUse: tracevault")); +} + +#[test] +fn claude_code_parse_system_record() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("claude-code"); + let chunk = json!({"type": "system", "subtype": "turn_duration", "durationMs": 5000.0}); + let record = adapter.parse_transcript_record(&chunk).unwrap(); + assert_eq!(record.record_type, "system"); + assert!(record.text.as_ref().unwrap().contains("5.0s")); +} + +#[test] +fn codex_map_event_types() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("codex"); + assert!(matches!( + adapter.map_event_type("SessionStart"), + StreamEventType::SessionStart + )); + assert!(matches!( + adapter.map_event_type("Stop"), + StreamEventType::SessionEnd + )); + assert!(matches!( + adapter.map_event_type("PostToolUse"), + StreamEventType::ToolUse + )); +} + +#[test] +fn codex_extract_token_usage() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("codex"); + let chunk = json!({"type": "event_msg", "payload": {"type": "token_count", "info": {"last_token_usage": {"input_tokens": 2000, "output_tokens": 300, "cached_input_tokens": 1500}}}}); + let usage = adapter.extract_token_usage(&chunk).unwrap(); + assert_eq!(usage.input_tokens, 2000); + assert_eq!(usage.output_tokens, 300); + assert_eq!(usage.cache_read_tokens, 1500); + assert_eq!(usage.cache_write_tokens, 0); +} + +#[test] +fn codex_extract_token_usage_non_token_chunk_returns_none() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("codex"); + let chunk = json!({"type": "event_msg", "payload": {"type": "agent_message"}}); + assert!(adapter.extract_token_usage(&chunk).is_none()); +} + +#[test] +fn codex_extract_model() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("codex"); + let chunk = json!({"type": "turn_context", "payload": {"model": "codex-mini-latest"}}); + assert_eq!( + adapter.extract_model(&chunk).as_deref(), + Some("codex-mini-latest") + ); +} + +#[test] +fn codex_extract_model_non_turn_context_returns_none() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("codex"); + let chunk = json!({"type": "event_msg", "payload": {"type": "agent_message"}}); + assert!(adapter.extract_model(&chunk).is_none()); +} + +#[test] +fn codex_parse_agent_message() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("codex"); + let chunk = json!({"type": "event_msg", "payload": {"type": "agent_message", "content": "I'll fix that bug now."}}); + let record = adapter.parse_transcript_record(&chunk).unwrap(); + assert_eq!(record.record_type, "assistant"); + assert_eq!(record.text.as_deref(), Some("I'll fix that bug now.")); +} + +#[test] +fn codex_parse_user_message() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("codex"); + let chunk = json!({"type": "event_msg", "payload": {"type": "user_message", "content": "Fix the login bug"}}); + let record = adapter.parse_transcript_record(&chunk).unwrap(); + assert_eq!(record.record_type, "user"); + assert_eq!(record.text.as_deref(), Some("Fix the login bug")); +} + +#[test] +fn codex_parse_shell_call() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("codex"); + let chunk = json!({"type": "response_item", "payload": {"type": "local_shell_call", "command": "cargo test", "output": "test result: ok. 5 passed"}}); + let record = adapter.parse_transcript_record(&chunk).unwrap(); + assert_eq!(record.record_type, "assistant"); + assert_eq!(record.tool_name.as_deref(), Some("Bash")); + assert!(record.text.as_ref().unwrap().contains("cargo test")); +} + +#[test] +fn codex_parse_token_count_returns_none_for_display() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("codex"); + let chunk = json!({"type": "event_msg", "payload": {"type": "token_count", "info": {"last_token_usage": {"input_tokens": 100, "output_tokens": 50}}}}); + assert!(adapter.parse_transcript_record(&chunk).is_none()); +} + +// Codex file changes are extracted from transcript, not hook events. +// These tests use parse_codex_patch directly. + +#[test] +fn codex_patch_parse_add_file() { + let changes = tracevault_core::agent_adapter::codex::parse_codex_patch( + "*** Begin Patch\n*** Add File: src/new.rs\n+fn main() {}\n*** End Patch\n", + ); + assert_eq!(changes.len(), 1); + assert_eq!(changes[0].file_path, "src/new.rs"); + assert_eq!(changes[0].change_type, "create"); + assert!(changes[0].content_hash.is_some()); +} + +#[test] +fn codex_patch_parse_update_file() { + let changes = tracevault_core::agent_adapter::codex::parse_codex_patch( + "*** Begin Patch\n*** Update File: src/lib.rs\n@@ fn old()\n-fn old()\n+fn new_func()\n*** End Patch\n", + ); + assert_eq!(changes.len(), 1); + assert_eq!(changes[0].file_path, "src/lib.rs"); + assert_eq!(changes[0].change_type, "edit"); + assert!(changes[0].diff_text.is_some()); +} + +#[test] +fn codex_patch_parse_delete_file() { + let changes = tracevault_core::agent_adapter::codex::parse_codex_patch( + "*** Begin Patch\n*** Delete File: src/old.rs\n*** End Patch\n", + ); + assert_eq!(changes.len(), 1); + assert_eq!(changes[0].file_path, "src/old.rs"); + assert_eq!(changes[0].change_type, "delete"); +} + +#[test] +fn codex_hook_extract_file_changes_returns_empty() { + // Codex does not extract file changes from hook events + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("codex"); + let input = json!({"command": "cargo build"}); + assert!(adapter.extract_file_changes("Bash", &input).is_empty()); +} + +#[test] +fn codex_is_file_modifying_always_false() { + // Codex file changes come from transcript, not hook events + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("codex"); + assert!(!adapter.is_file_modifying("Bash")); + assert!(!adapter.is_file_modifying("Read")); + assert!(!adapter.is_file_modifying("apply_patch")); +} + +#[test] +fn codex_extract_file_changes_from_transcript_apply_patch() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("codex"); + let chunk = json!({ + "type": "response_item", + "payload": { + "type": "custom_tool_call", + "name": "apply_patch", + "input": "*** Begin Patch\n*** Update File: src/main.rs\n@@ fn old()\n-fn old()\n+fn new_func()\n*** End Patch\n" + } + }); + let changes = adapter.extract_file_changes_from_transcript(&chunk); + assert_eq!(changes.len(), 1); + assert_eq!(changes[0].file_path, "src/main.rs"); + assert_eq!(changes[0].change_type, "edit"); +} + +#[test] +fn codex_extract_file_changes_from_transcript_non_patch_returns_empty() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("codex"); + let chunk = json!({ + "type": "response_item", + "payload": {"type": "message", "role": "assistant", "content": []} + }); + assert!(adapter + .extract_file_changes_from_transcript(&chunk) + .is_empty()); +} + +#[test] +fn codex_reasoning_record_returns_none() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("codex"); + let chunk = json!({ + "type": "response_item", + "payload": { + "type": "reasoning", + "content": null, + "summary": [], + "encrypted_content": "gAAAAA..." + } + }); + assert!(adapter.parse_transcript_record(&chunk).is_none()); +} + +#[test] +fn codex_custom_tool_call_display() { + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("codex"); + let chunk = json!({ + "type": "response_item", + "timestamp": "2026-04-03T17:52:42Z", + "payload": { + "type": "custom_tool_call", + "name": "apply_patch", + "input": "*** Begin Patch\n*** Update File: README.md\n@@\n old line\n+new line\n*** End Patch" + } + }); + let record = adapter.parse_transcript_record(&chunk).unwrap(); + assert_eq!(record.record_type, "assistant"); + assert_eq!(record.tool_name.as_deref(), Some("apply_patch")); + assert!(record.text.as_ref().unwrap().contains("Update File")); +} + +#[test] +fn claude_code_extract_file_changes_from_transcript_returns_empty() { + // Claude Code file changes come from hook events, not transcript + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("claude-code"); + let chunk = json!({"type": "assistant", "message": {"content": []}}); + assert!(adapter + .extract_file_changes_from_transcript(&chunk) + .is_empty()); +} diff --git a/crates/tracevault-core/tests/streaming_test.rs b/crates/tracevault-core/tests/streaming_test.rs index 19f3f2b..677a2b2 100644 --- a/crates/tracevault-core/tests/streaming_test.rs +++ b/crates/tracevault-core/tests/streaming_test.rs @@ -28,50 +28,6 @@ fn test_stream_event_request_serialization() { assert_eq!(parsed.event_index, Some(42)); } -#[test] -fn test_extract_file_change_from_edit() { - let tool_input = json!({ - "file_path": "/repo/src/lib.rs", - "old_string": "fn old() {}", - "new_string": "fn new_func() {}" - }); - let change = extract_file_change("Edit", &tool_input); - assert!(change.is_some()); - let c = change.unwrap(); - assert_eq!(c.file_path, "/repo/src/lib.rs"); - assert_eq!(c.change_type, "edit"); - assert!(c.diff_text.is_some()); -} - -#[test] -fn test_extract_file_change_from_write() { - let tool_input = json!({ - "file_path": "/repo/src/new_file.rs", - "content": "fn main() {}" - }); - let change = extract_file_change("Write", &tool_input); - assert!(change.is_some()); - let c = change.unwrap(); - assert_eq!(c.file_path, "/repo/src/new_file.rs"); - assert_eq!(c.change_type, "create"); - assert!(c.content_hash.is_some()); -} - -#[test] -fn test_extract_file_change_from_read_returns_none() { - let tool_input = json!({"file_path": "/repo/src/lib.rs"}); - assert!(extract_file_change("Read", &tool_input).is_none()); -} - -#[test] -fn test_is_file_modifying_tool() { - assert!(is_file_modifying_tool("Write")); - assert!(is_file_modifying_tool("Edit")); - assert!(is_file_modifying_tool("Bash")); - assert!(!is_file_modifying_tool("Read")); - assert!(!is_file_modifying_tool("Grep")); -} - #[test] fn test_commit_push_request_serialization() { let req = CommitPushRequest { @@ -86,21 +42,3 @@ fn test_commit_push_request_serialization() { let parsed: CommitPushRequest = serde_json::from_str(&json_str).unwrap(); assert_eq!(parsed.commit_sha, "abc123"); } - -#[test] -fn extract_file_change_write_missing_content() { - let input = json!({"file_path": "/tmp/test.rs"}); - assert!(extract_file_change("Write", &input).is_none()); -} - -#[test] -fn extract_file_change_edit_missing_old_string() { - let input = json!({"file_path": "/tmp/test.rs", "new_string": "new"}); - assert!(extract_file_change("Edit", &input).is_none()); -} - -#[test] -fn extract_file_change_write_missing_file_path() { - let input = json!({"content": "hello"}); - assert!(extract_file_change("Write", &input).is_none()); -} diff --git a/crates/tracevault-server/src/api/session_detail.rs b/crates/tracevault-server/src/api/session_detail.rs index 57994ed..f06b2d6 100644 --- a/crates/tracevault-server/src/api/session_detail.rs +++ b/crates/tracevault-server/src/api/session_detail.rs @@ -4,6 +4,8 @@ use chrono::{DateTime, Utc}; use serde::Serialize; use uuid::Uuid; +use tracevault_core::agent_adapter::AgentAdapter; + use crate::error::AppError; use crate::extractors::OrgAuth; use crate::pricing::{self, ModelPricing}; @@ -91,230 +93,10 @@ pub struct TranscriptRecord { pub model: Option, } -fn parse_record(record: &serde_json::Value, pricing: &ModelPricing) -> Option { - let record_type = record.get("type")?.as_str()?.to_string(); - let timestamp = record - .get("timestamp") - .and_then(|v| v.as_str()) - .map(String::from); - - match record_type.as_str() { - "assistant" => { - let msg = match record.get("message") { - Some(m) => m, - None => { - return Some(TranscriptRecord { - record_type, - timestamp, - content_types: vec![], - tool_name: None, - text: None, - usage: None, - model: None, - }); - } - }; - let model = msg.get("model").and_then(|v| v.as_str()).map(String::from); - - let mut content_types = Vec::new(); - let mut texts = Vec::new(); - if let Some(content) = msg.get("content").and_then(|v| v.as_array()) { - for block in content { - if let Some(ct) = block.get("type").and_then(|v| v.as_str()) { - if !content_types.contains(&ct.to_string()) { - content_types.push(ct.to_string()); - } - } - if let Some(text) = block.get("text").and_then(|v| v.as_str()) { - texts.push(text.to_string()); - } - if let Some(thinking) = block.get("thinking").and_then(|v| v.as_str()) { - texts.push(format!("[thinking] {}", thinking)); - } - } - } - - let usage = msg.get("usage").map(|u| { - let total_input = u.get("input_tokens").and_then(|v| v.as_i64()).unwrap_or(0); - let output = u.get("output_tokens").and_then(|v| v.as_i64()).unwrap_or(0); - let cache_read = u - .get("cache_read_input_tokens") - .and_then(|v| v.as_i64()) - .unwrap_or(0); - let cache_write = u - .get("cache_creation_input_tokens") - .and_then(|v| v.as_i64()) - .unwrap_or(0); - // input_tokens from the API includes cache_read and cache_write tokens, - // so subtract them to get fresh (non-cached) input tokens only - let fresh_input = (total_input - cache_read - cache_write).max(0); - let cost = pricing::estimate_cost_with_pricing( - pricing, - fresh_input, - output, - cache_read, - cache_write, - ); - RecordUsage { - input_tokens: fresh_input, - output_tokens: output, - cache_read_tokens: cache_read, - cache_write_tokens: cache_write, - cost_usd: cost, - } - }); - - let tool_name = msg - .get("content") - .and_then(|v| v.as_array()) - .and_then(|arr| { - arr.iter() - .find(|b| b.get("type").and_then(|v| v.as_str()) == Some("tool_use")) - }) - .and_then(|b| b.get("name").and_then(|v| v.as_str()).map(String::from)); - - Some(TranscriptRecord { - record_type, - timestamp, - content_types, - tool_name, - text: if texts.is_empty() { - None - } else { - Some(texts.join("\n\n")) - }, - usage, - model, - }) - } - "user" => { - let mut content_types = Vec::new(); - let mut text = None; - let mut tool_name = None; - - let msg = record.get("message"); - match msg.and_then(|m| m.get("content")) { - Some(serde_json::Value::String(s)) => { - content_types.push("text".to_string()); - text = Some(s.clone()); - } - Some(serde_json::Value::Array(arr)) => { - for block in arr { - if let Some(ct) = block.get("type").and_then(|v| v.as_str()) { - if !content_types.contains(&ct.to_string()) { - content_types.push(ct.to_string()); - } - if ct == "tool_result" { - if let Some(content) = block.get("content").and_then(|v| v.as_str()) - { - text = Some(content.to_string()); - } - } else if ct == "text" { - if let Some(t) = block.get("text").and_then(|v| v.as_str()) { - text = Some(t.to_string()); - } - } - } - } - } - _ => {} - } - - if let Some(tur) = record.get("toolUseResult") { - if let Some(file) = tur - .get("file") - .and_then(|f| f.get("filePath").and_then(|v| v.as_str())) - { - tool_name = Some(format!("Read: {}", file)); - } else if tur.get("filenames").is_some() { - tool_name = Some("Glob".to_string()); - } else if tur.get("stdout").is_some() { - tool_name = Some("Bash".to_string()); - } - } - - Some(TranscriptRecord { - record_type, - timestamp, - content_types, - tool_name, - text, - usage: None, - model: None, - }) - } - "progress" => { - let data = record.get("data"); - let hook_name = data - .and_then(|d| d.get("hookName").and_then(|v| v.as_str())) - .map(String::from); - let hook_event = data.and_then(|d| d.get("hookEvent").and_then(|v| v.as_str())); - let text = hook_event.map(|e| { - if let Some(ref name) = hook_name { - format!("{}: {}", e, name) - } else { - e.to_string() - } - }); - - Some(TranscriptRecord { - record_type, - timestamp, - content_types: vec![], - tool_name: hook_name, - text, - usage: None, - model: None, - }) - } - "system" => { - let subtype = record - .get("subtype") - .and_then(|v| v.as_str()) - .unwrap_or("unknown"); - let text = match subtype { - "turn_duration" => { - let ms = record - .get("durationMs") - .and_then(|v| v.as_f64()) - .unwrap_or(0.0); - Some(format!("turn_duration: {:.1}s", ms / 1000.0)) - } - "stop_hook_summary" => { - let count = record - .get("hookCount") - .and_then(|v| v.as_i64()) - .unwrap_or(0); - Some(format!("stop_hook_summary: {} hooks", count)) - } - _ => Some(subtype.to_string()), - }; - - Some(TranscriptRecord { - record_type, - timestamp, - content_types: vec![subtype.to_string()], - tool_name: None, - text, - usage: None, - model: None, - }) - } - _ => Some(TranscriptRecord { - record_type, - timestamp, - content_types: vec![], - tool_name: None, - text: None, - usage: None, - model: None, - }), - } -} - pub fn parse_transcript( transcript: &serde_json::Value, pricing: &ModelPricing, + adapter: &dyn AgentAdapter, ) -> ( Vec, Vec, @@ -335,7 +117,41 @@ pub fn parse_transcript( let mut total_cache_write: i64 = 0; for record in records { - if let Some(tr) = parse_record(record, pricing) { + if let Some(parsed) = adapter.parse_transcript_record(record) { + let usage = if parsed.raw_input_tokens.is_some() { + let total_input_raw = parsed.raw_input_tokens.unwrap_or(0); + let output = parsed.raw_output_tokens.unwrap_or(0); + let cache_read = parsed.raw_cache_read_tokens.unwrap_or(0); + let cache_write = parsed.raw_cache_write_tokens.unwrap_or(0); + let fresh_input = (total_input_raw - cache_read - cache_write).max(0); + let cost = pricing::estimate_cost_with_pricing( + pricing, + fresh_input, + output, + cache_read, + cache_write, + ); + Some(RecordUsage { + input_tokens: fresh_input, + output_tokens: output, + cache_read_tokens: cache_read, + cache_write_tokens: cache_write, + cost_usd: cost, + }) + } else { + None + }; + + let tr = TranscriptRecord { + record_type: parsed.record_type.clone(), + timestamp: parsed.timestamp, + content_types: parsed.content_types, + tool_name: parsed.tool_name, + text: parsed.text, + usage, + model: parsed.model, + }; + if tr.record_type == "assistant" { if let Some(ref usage) = tr.usage { let model = tr.model.as_deref().unwrap_or("unknown"); @@ -414,6 +230,7 @@ pub fn parse_transcript( struct SessionRow { session_id: String, model: Option, + tool: Option, started_at: Option>, ended_at: Option>, duration_ms: Option, @@ -436,7 +253,7 @@ pub async fn get_session_detail( let org_id = auth.org_id; let row = sqlx::query_as::<_, SessionRow>( - "SELECT s.session_id, s.model, s.started_at, s.ended_at, s.duration_ms, + "SELECT s.session_id, s.model, s.tool, s.started_at, s.ended_at, s.duration_ms, s.total_tokens, s.input_tokens, s.output_tokens, s.cache_read_tokens, s.cache_write_tokens, s.estimated_cost_usd, @@ -471,8 +288,11 @@ pub async fn get_session_detail( let transcript_array: Vec = chunks.into_iter().map(|(d,)| d).collect(); let transcript_val = serde_json::Value::Array(transcript_array); + let adapter = state + .agent_registry + .get(row.tool.as_deref().unwrap_or("claude-code")); let (per_call, transcript_records, token_distribution, cost_breakdown, cache_savings) = - parse_transcript(&transcript_val, &pricing); + parse_transcript(&transcript_val, &pricing, adapter); // Count API calls from per_call data since api_calls column doesn't exist on sessions let api_calls = per_call.len() as i32; @@ -505,6 +325,7 @@ pub async fn get_session_detail( #[cfg(test)] mod tests { use super::*; + use tracevault_core::agent_adapter::AgentAdapterRegistry; fn test_pricing() -> ModelPricing { ModelPricing { @@ -518,8 +339,10 @@ mod tests { #[test] fn test_parse_empty_transcript() { let transcript = serde_json::json!([]); + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("claude-code"); let (per_call, records, dist, cost, savings) = - parse_transcript(&transcript, &test_pricing()); + parse_transcript(&transcript, &test_pricing(), adapter); assert!(per_call.is_empty()); assert!(records.is_empty()); assert_eq!(dist.input_tokens, 0); @@ -545,8 +368,10 @@ mod tests { } } ]); + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("claude-code"); let (per_call, records, dist, _cost, _savings) = - parse_transcript(&transcript, &test_pricing()); + parse_transcript(&transcript, &test_pricing(), adapter); assert_eq!(per_call.len(), 1); assert_eq!(per_call[0].index, 1); assert_eq!(per_call[0].cache_read_tokens, 1000); @@ -567,8 +392,10 @@ mod tests { } } ]); + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("claude-code"); let (per_call, records, _dist, _cost, _savings) = - parse_transcript(&transcript, &test_pricing()); + parse_transcript(&transcript, &test_pricing(), adapter); assert!(per_call.is_empty()); assert_eq!(records.len(), 1); } @@ -592,7 +419,10 @@ mod tests { } ]); let pricing = test_pricing(); - let (_per_call, _records, _dist, _cost, savings) = parse_transcript(&transcript, &pricing); + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("claude-code"); + let (_per_call, _records, _dist, _cost, savings) = + parse_transcript(&transcript, &pricing, adapter); assert!((savings.gross_savings_usd - 13.5).abs() < 0.001); assert!((savings.cache_write_overhead_usd - 0.375).abs() < 0.001); assert!((savings.net_savings_usd - 13.125).abs() < 0.001); @@ -620,7 +450,9 @@ mod tests { } } ]); - let (per_call, _, _, _, _) = parse_transcript(&transcript, &test_pricing()); + let registry = AgentAdapterRegistry::new(); + let adapter = registry.get("claude-code"); + let (per_call, _, _, _, _) = parse_transcript(&transcript, &test_pricing(), adapter); assert_eq!(per_call.len(), 2); assert!((per_call[0].cumulative_cost_usd - 15.0).abs() < 0.001); assert!((per_call[1].cumulative_cost_usd - 30.0).abs() < 0.001); diff --git a/crates/tracevault-server/src/api/traces_ui.rs b/crates/tracevault-server/src/api/traces_ui.rs index 9ec4742..9331ed7 100644 --- a/crates/tracevault-server/src/api/traces_ui.rs +++ b/crates/tracevault-server/src/api/traces_ui.rs @@ -438,17 +438,14 @@ pub async fn get_session_transcript( ) -> Result, AppError> { verify_session_access(&state.pool, session_id, auth.org_id).await?; - let session_model: Option = - sqlx::query_scalar("SELECT model FROM sessions WHERE id = $1") - .bind(session_id) - .fetch_one(&state.pool) - .await?; - - let session_started_at: Option> = - sqlx::query_scalar("SELECT started_at FROM sessions WHERE id = $1") - .bind(session_id) - .fetch_one(&state.pool) - .await?; + let (session_model, session_tool, session_started_at): ( + Option, + Option, + Option>, + ) = sqlx::query_as("SELECT model, tool, started_at FROM sessions WHERE id = $1") + .bind(session_id) + .fetch_one(&state.pool) + .await?; let transcript_chunks = sqlx::query_as::<_, TranscriptChunkRow>( "SELECT chunk_index, data @@ -467,10 +464,13 @@ pub async fn get_session_transcript( ) .await; + let adapter = state + .agent_registry + .get(session_tool.as_deref().unwrap_or("claude-code")); let transcript_array: Vec = transcript_chunks.iter().map(|c| c.data.clone()).collect(); let transcript_val = serde_json::Value::Array(transcript_array); - let (_, transcript_records, _, _, _) = parse_transcript(&transcript_val, &pricing); + let (_, transcript_records, _, _, _) = parse_transcript(&transcript_val, &pricing, adapter); Ok(Json(TranscriptResponse { transcript_chunks, diff --git a/crates/tracevault-server/src/lib.rs b/crates/tracevault-server/src/lib.rs index 518ddd0..eea7d8d 100644 --- a/crates/tracevault-server/src/lib.rs +++ b/crates/tracevault-server/src/lib.rs @@ -22,6 +22,9 @@ pub mod story; pub use error::AppError; +use std::sync::Arc; +use tracevault_core::agent_adapter::AgentAdapterRegistry; + #[derive(Clone)] pub struct AppState { pub pool: sqlx::PgPool, @@ -31,4 +34,5 @@ pub struct AppState { pub http_client: reqwest::Client, pub cors_origin: String, pub invite_expiry_minutes: u64, + pub agent_registry: Arc, } diff --git a/crates/tracevault-server/src/main.rs b/crates/tracevault-server/src/main.rs index b32a16d..e0afebf 100644 --- a/crates/tracevault-server/src/main.rs +++ b/crates/tracevault-server/src/main.rs @@ -9,6 +9,7 @@ use tower_governor::{governor::GovernorConfigBuilder, GovernorLayer}; use tower_http::cors::CorsLayer; use tower_http::trace::TraceLayer; +use tracevault_core::agent_adapter::AgentAdapterRegistry; use tracevault_server::{api, config, db, extensions, pricing_sync, repo_manager, AppState}; #[tokio::main] @@ -495,6 +496,7 @@ async fn main() { http_client: http_client.clone(), cors_origin: cfg.cors_origin.clone(), invite_expiry_minutes: cfg.invite_expiry_minutes, + agent_registry: Arc::new(AgentAdapterRegistry::new()), }); let listener = tokio::net::TcpListener::bind(&bind_addr).await.unwrap(); diff --git a/crates/tracevault-server/src/service/stream.rs b/crates/tracevault-server/src/service/stream.rs index 592902b..a3fffcf 100644 --- a/crates/tracevault-server/src/service/stream.rs +++ b/crates/tracevault-server/src/service/stream.rs @@ -1,8 +1,5 @@ use tracevault_core::software::extract_software; -use tracevault_core::streaming::{ - extract_file_change, is_file_modifying_tool, StreamEventRequest, StreamEventResponse, - StreamEventType, -}; +use tracevault_core::streaming::{StreamEventRequest, StreamEventResponse, StreamEventType}; use uuid::Uuid; use crate::error::AppError; @@ -37,6 +34,9 @@ impl StreamService { Some("claude-code".to_string()) }; + let agent_name = tool.as_deref().unwrap_or("claude-code"); + let adapter = state.agent_registry.get(agent_name); + // 3. Upsert session let session_db_id = SessionRepo::upsert( &state.pool, @@ -83,31 +83,55 @@ impl StreamService { continue; } - // Extract token usage from assistant messages - if let Some(msg) = line.get("message") { - if let Some(usage) = msg.get("usage") { - batch_input += usage - .get("input_tokens") - .and_then(|v| v.as_i64()) - .unwrap_or(0); - batch_output += usage - .get("output_tokens") - .and_then(|v| v.as_i64()) - .unwrap_or(0); - batch_cache_read += usage - .get("cache_read_input_tokens") - .and_then(|v| v.as_i64()) - .unwrap_or(0); - batch_cache_write += usage - .get("cache_creation_input_tokens") - .and_then(|v| v.as_i64()) - .unwrap_or(0); - } - if detected_model.is_none() { - detected_model = - msg.get("model").and_then(|v| v.as_str()).map(String::from); + // Extract file changes from transcript chunks (e.g. Codex apply_patch). + // Each adapter decides which chunk types contain file modifications. + let transcript_file_changes = + adapter.extract_file_changes_from_transcript(line); + for change in transcript_file_changes { + let tool_name = line + .get("payload") + .and_then(|p| p.get("name")) + .and_then(|v| v.as_str()) + .unwrap_or("unknown"); + let event_id = EventRepo::insert_tool_event( + &state.pool, + &crate::repo::events::InsertToolEvent { + session_id: session_db_id, + event_index: chunk_index, + tool_name: Some(tool_name.to_string()), + tool_input: line.get("payload").cloned(), + tool_response: None, + timestamp: Some(req.timestamp), + }, + ) + .await?; + if let Some(eid) = event_id { + EventRepo::insert_file_change( + &state.pool, + &InsertFileChange { + session_id: session_db_id, + event_id: eid, + file_path: change.file_path, + change_type: change.change_type, + diff_text: change.diff_text, + content_hash: change.content_hash, + timestamp: Some(req.timestamp), + }, + ) + .await?; } } + + // Extract token usage via adapter + if let Some(usage) = adapter.extract_token_usage(line) { + batch_input += usage.input_tokens; + batch_output += usage.output_tokens; + batch_cache_read += usage.cache_read_tokens; + batch_cache_write += usage.cache_write_tokens; + } + if detected_model.is_none() { + detected_model = adapter.extract_model(line); + } } // Update session token counts and cost if we found usage data @@ -156,7 +180,7 @@ impl StreamService { })?; let tool_name = req.tool_name.as_deref().unwrap_or(""); - let store_response = is_file_modifying_tool(tool_name); + let store_response = adapter.is_file_modifying(tool_name); let inserted_id = EventRepo::insert_tool_event( &state.pool, @@ -179,9 +203,10 @@ impl StreamService { event_db_id = Some(eid); // Extract file changes for file-modifying tools - if is_file_modifying_tool(tool_name) { + if adapter.is_file_modifying(tool_name) { if let Some(ref tool_input) = req.tool_input { - if let Some(change) = extract_file_change(tool_name, tool_input) { + let file_changes = adapter.extract_file_changes(tool_name, tool_input); + for change in file_changes { EventRepo::insert_file_change( &state.pool, &InsertFileChange { diff --git a/web/src/lib/components/AgentBadge.svelte b/web/src/lib/components/AgentBadge.svelte new file mode 100644 index 0000000..e283fca --- /dev/null +++ b/web/src/lib/components/AgentBadge.svelte @@ -0,0 +1,69 @@ + + +{#if agent} + + {#if tool === 'claude-code'} + + + + + {:else if tool === 'codex'} + + + + + + {:else if tool === 'gemini'} + + + + + {:else if tool === 'cursor'} + + + + + {:else} + + + + + + + + {/if} + {agent.label} + +{/if} diff --git a/web/src/routes/orgs/[slug]/traces/sessions/+page.svelte b/web/src/routes/orgs/[slug]/traces/sessions/+page.svelte index 3ff5daa..bfe0cb2 100644 --- a/web/src/routes/orgs/[slug]/traces/sessions/+page.svelte +++ b/web/src/routes/orgs/[slug]/traces/sessions/+page.svelte @@ -6,6 +6,7 @@ import type { SessionItem } from '$lib/types'; import DataTable from '$lib/components/DataTable.svelte'; import StatusBadge from '$lib/components/StatusBadge.svelte'; + import AgentBadge from '$lib/components/AgentBadge.svelte'; import LoadingState from '$lib/components/LoadingState.svelte'; import ErrorState from '$lib/components/ErrorState.svelte'; import EmptyState from '$lib/components/EmptyState.svelte'; @@ -87,7 +88,10 @@ > {#snippet children({ row, col })} {#if col.key === '_status'} - +
+ + +
{:else if col.key === 'session_id'} {String(row.session_id).slice(0, 8)} diff --git a/web/src/routes/orgs/[slug]/traces/sessions/[id]/+page.svelte b/web/src/routes/orgs/[slug]/traces/sessions/[id]/+page.svelte index b4ed381..39f854f 100644 --- a/web/src/routes/orgs/[slug]/traces/sessions/[id]/+page.svelte +++ b/web/src/routes/orgs/[slug]/traces/sessions/[id]/+page.svelte @@ -15,6 +15,7 @@ import { formatDateTime } from '$lib/utils/date'; import * as Table from '$lib/components/ui/table/index.js'; import StatusBadge from '$lib/components/StatusBadge.svelte'; + import AgentBadge from '$lib/components/AgentBadge.svelte'; import LoadingState from '$lib/components/LoadingState.svelte'; import ErrorState from '$lib/components/ErrorState.svelte'; import SessionTranscript from '$lib/components/session-detail/SessionTranscript.svelte'; @@ -197,6 +198,7 @@ / {session.session_id.slice(0, 8)} + From 1ab79fd1cebe208fb6981da9f9bf490091fa76a2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Grabowski?= <> Date: Fri, 3 Apr 2026 21:09:44 +0200 Subject: [PATCH 2/5] refactor: replace --codex flag with repeatable --agent option Allow multiple agents via --agent flag (e.g. --agent codex --agent gemini) instead of single --codex bool. Extensible for future agents. Co-Authored-By: Claude Opus 4.6 (1M context) --- README.md | 6 +++--- crates/tracevault-cli/src/commands/init.rs | 11 +++++++---- crates/tracevault-cli/src/main.rs | 14 +++++++------- 3 files changed, 17 insertions(+), 14 deletions(-) diff --git a/README.md b/README.md index bf45e2d..3b75196 100644 --- a/README.md +++ b/README.md @@ -268,13 +268,13 @@ That's it. From this point on, every Claude Code session in this repo is automat ## Using with Codex CLI -[Codex CLI](https://github.com/openai/codex) (OpenAI's coding agent) is also supported. Initialize with the `--codex` flag to install Codex hooks: +[Codex CLI](https://github.com/openai/codex) (OpenAI's coding agent) is also supported. Initialize with the `--agent codex` flag to install Codex hooks: ```sh npm install -g @openai/codex cd /path/to/your/repo tracevault login --server-url https://your-tracevault-server.example.com -tracevault init --codex +tracevault init --agent codex ``` This installs hooks in `.codex/hooks.json` in addition to the Claude Code hooks. Codex sessions are traced including transcript parsing, token usage, and file changes via `apply_patch`. The session detail view shows a Codex badge to distinguish agent types. @@ -331,7 +331,7 @@ export DATABASE_URL=postgres://user:password@host:5432/tracevault?sslmode=requir | Command | Description | |---------|-------------| -| `tracevault init [--server-url URL] [--codex]` | Initialize TraceVault in current repo, install hooks (Claude Code by default, `--codex` adds Codex CLI) | +| `tracevault init [--server-url URL] [--agent ]` | Initialize TraceVault in current repo, install hooks (Claude Code by default, `--agent` adds extra agents e.g. `codex`) | | `tracevault login --server-url URL` | Authenticate via device auth flow (opens browser) | | `tracevault logout` | Clear local credentials | | `tracevault hook --event ` | Handle a hook event from any agent (reads JSON from stdin) | diff --git a/crates/tracevault-cli/src/commands/init.rs b/crates/tracevault-cli/src/commands/init.rs index 104229d..a06ce4a 100644 --- a/crates/tracevault-cli/src/commands/init.rs +++ b/crates/tracevault-cli/src/commands/init.rs @@ -33,7 +33,7 @@ fn parse_github_org(remote_url: &str) -> Option { pub async fn init_in_directory( project_root: &Path, server_url: Option<&str>, - codex: bool, + agents: &[String], ) -> Result<(), io::Error> { // Check for git repository if !project_root.join(".git").exists() { @@ -79,9 +79,12 @@ pub async fn init_in_directory( // Install Claude Code hooks into .claude/settings.json install_claude_hooks(project_root)?; - // Install Codex hooks into .codex/hooks.json - if codex { - install_codex_hooks(project_root)?; + // Install agent-specific hooks + for agent in agents { + match agent.as_str() { + "codex" => install_codex_hooks(project_root)?, + other => eprintln!("Warning: unknown agent '{}', skipping hooks", other), + } } // Install git hooks diff --git a/crates/tracevault-cli/src/main.rs b/crates/tracevault-cli/src/main.rs index bf18460..81ababc 100644 --- a/crates/tracevault-cli/src/main.rs +++ b/crates/tracevault-cli/src/main.rs @@ -15,9 +15,9 @@ enum Cli { /// TraceVault server URL for repo registration #[arg(long)] server_url: Option, - /// Also install Codex CLI hooks (.codex/hooks.json) - #[arg(long)] - codex: bool, + /// Additional AI agents to install hooks for (e.g. codex, gemini) + #[arg(long = "agent")] + agents: Vec, }, /// Show current session status Status, @@ -69,14 +69,14 @@ enum Cli { async fn main() { let cli = Cli::parse(); match cli { - Cli::Init { server_url, codex } => { + Cli::Init { server_url, agents } => { let cwd = env::current_dir().expect("Cannot determine current directory"); - match commands::init::init_in_directory(&cwd, server_url.as_deref(), codex).await { + match commands::init::init_in_directory(&cwd, server_url.as_deref(), &agents).await { Ok(()) => { println!("TraceVault initialized in {}", cwd.display()); println!("Claude Code hooks installed in .claude/settings.json"); - if codex { - println!("Codex hooks installed in .codex/hooks.json"); + for agent in &agents { + println!("{} hooks installed", agent); } println!("Git pre-push hook installed"); } From 0268f2585309d7d6db73218be5fe6fc1dfd6f806 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Grabowski?= <> Date: Fri, 3 Apr 2026 21:18:37 +0200 Subject: [PATCH 3/5] fix: save model name even when transcript batch has no tokens Codex sends model in turn_context chunks which may arrive in a separate batch from token_count chunks. Previously model was only saved when the batch contained tokens. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/tracevault-server/src/service/stream.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/tracevault-server/src/service/stream.rs b/crates/tracevault-server/src/service/stream.rs index a3fffcf..7bbd4ba 100644 --- a/crates/tracevault-server/src/service/stream.rs +++ b/crates/tracevault-server/src/service/stream.rs @@ -139,7 +139,7 @@ impl StreamService { || batch_output > 0 || batch_cache_read > 0 || batch_cache_write > 0; - if has_tokens { + if has_tokens || detected_model.is_some() { let model_name = detected_model.as_deref().unwrap_or("unknown"); // input_tokens from the API includes cache_read and cache_write, // subtract to get fresh (non-cached) input only From f13b7acd01e318e1711a653f87ecb238d9e0c81d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Grabowski?= <> Date: Tue, 7 Apr 2026 19:01:22 +0200 Subject: [PATCH 4/5] fix: remove unused TokenUsage import in agent adapter test --- crates/tracevault-core/tests/agent_adapter_test.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/tracevault-core/tests/agent_adapter_test.rs b/crates/tracevault-core/tests/agent_adapter_test.rs index 2cc99dd..57dd430 100644 --- a/crates/tracevault-core/tests/agent_adapter_test.rs +++ b/crates/tracevault-core/tests/agent_adapter_test.rs @@ -1,5 +1,5 @@ use serde_json::json; -use tracevault_core::agent_adapter::{AgentAdapterRegistry, TokenUsage}; +use tracevault_core::agent_adapter::AgentAdapterRegistry; use tracevault_core::streaming::StreamEventType; #[test] From c6136356801019478d2ccb9246d5c1f267239d91 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Grabowski?= <> Date: Tue, 7 Apr 2026 19:09:12 +0200 Subject: [PATCH 5/5] fix: make agents parameter optional in init_in_directory, default to claude --- crates/tracevault-cli/src/commands/init.rs | 10 +++---- crates/tracevault-cli/src/main.rs | 12 +++++++- crates/tracevault-cli/tests/e2e_test.rs | 4 +-- crates/tracevault-cli/tests/init_test.rs | 34 ++++++++++++---------- 4 files changed, 37 insertions(+), 23 deletions(-) diff --git a/crates/tracevault-cli/src/commands/init.rs b/crates/tracevault-cli/src/commands/init.rs index a06ce4a..e61b6c5 100644 --- a/crates/tracevault-cli/src/commands/init.rs +++ b/crates/tracevault-cli/src/commands/init.rs @@ -33,7 +33,7 @@ fn parse_github_org(remote_url: &str) -> Option { pub async fn init_in_directory( project_root: &Path, server_url: Option<&str>, - agents: &[String], + agents: Option<&[String]>, ) -> Result<(), io::Error> { // Check for git repository if !project_root.join(".git").exists() { @@ -76,12 +76,12 @@ pub async fn init_in_directory( "sessions/\ncache/\n*.local.toml\n", )?; - // Install Claude Code hooks into .claude/settings.json - install_claude_hooks(project_root)?; - - // Install agent-specific hooks + // Install agent-specific hooks (defaults to claude when none specified) + let default_agents = [String::from("claude")]; + let agents = agents.unwrap_or(&default_agents); for agent in agents { match agent.as_str() { + "claude" => install_claude_hooks(project_root)?, "codex" => install_codex_hooks(project_root)?, other => eprintln!("Warning: unknown agent '{}', skipping hooks", other), } diff --git a/crates/tracevault-cli/src/main.rs b/crates/tracevault-cli/src/main.rs index 81ababc..d3ab60c 100644 --- a/crates/tracevault-cli/src/main.rs +++ b/crates/tracevault-cli/src/main.rs @@ -71,7 +71,17 @@ async fn main() { match cli { Cli::Init { server_url, agents } => { let cwd = env::current_dir().expect("Cannot determine current directory"); - match commands::init::init_in_directory(&cwd, server_url.as_deref(), &agents).await { + match commands::init::init_in_directory( + &cwd, + server_url.as_deref(), + if agents.is_empty() { + None + } else { + Some(&agents) + }, + ) + .await + { Ok(()) => { println!("TraceVault initialized in {}", cwd.display()); println!("Claude Code hooks installed in .claude/settings.json"); diff --git a/crates/tracevault-cli/tests/e2e_test.rs b/crates/tracevault-cli/tests/e2e_test.rs index eecb772..b497d1e 100644 --- a/crates/tracevault-cli/tests/e2e_test.rs +++ b/crates/tracevault-cli/tests/e2e_test.rs @@ -12,7 +12,7 @@ async fn full_flow_init_hook_and_local_stats() { let tmp = tmp_git_repo(); // 1. Init - tracevault_cli::commands::init::init_in_directory(tmp.path(), None) + tracevault_cli::commands::init::init_in_directory(tmp.path(), None, None) .await .unwrap(); assert!(tmp.path().join(".tracevault/config.toml").exists()); @@ -91,7 +91,7 @@ async fn full_flow_init_hook_and_local_stats() { #[tokio::test] async fn multiple_sessions_tracked_independently() { let tmp = tmp_git_repo(); - tracevault_cli::commands::init::init_in_directory(tmp.path(), None) + tracevault_cli::commands::init::init_in_directory(tmp.path(), None, None) .await .unwrap(); diff --git a/crates/tracevault-cli/tests/init_test.rs b/crates/tracevault-cli/tests/init_test.rs index 3be8a39..3c0b4d8 100644 --- a/crates/tracevault-cli/tests/init_test.rs +++ b/crates/tracevault-cli/tests/init_test.rs @@ -10,7 +10,7 @@ fn tmp_git_repo() -> TempDir { #[tokio::test] async fn init_fails_without_git() { let tmp = TempDir::new().unwrap(); - let result = tracevault_cli::commands::init::init_in_directory(tmp.path(), None).await; + let result = tracevault_cli::commands::init::init_in_directory(tmp.path(), None, None).await; assert!(result.is_err()); assert!(result .unwrap_err() @@ -23,7 +23,7 @@ async fn init_creates_tracevault_config() { let tmp = tmp_git_repo(); let config_path = tmp.path().join(".tracevault").join("config.toml"); - tracevault_cli::commands::init::init_in_directory(tmp.path(), None) + tracevault_cli::commands::init::init_in_directory(tmp.path(), None, None) .await .unwrap(); @@ -36,7 +36,7 @@ async fn init_creates_tracevault_config() { async fn init_creates_directory_structure() { let tmp = tmp_git_repo(); - tracevault_cli::commands::init::init_in_directory(tmp.path(), None) + tracevault_cli::commands::init::init_in_directory(tmp.path(), None, None) .await .unwrap(); @@ -50,7 +50,7 @@ async fn init_creates_directory_structure() { async fn init_installs_claude_hooks() { let tmp = tmp_git_repo(); - tracevault_cli::commands::init::init_in_directory(tmp.path(), None) + tracevault_cli::commands::init::init_in_directory(tmp.path(), None, None) .await .unwrap(); @@ -74,7 +74,7 @@ async fn init_merges_into_existing_settings() { fs::create_dir_all(&claude_dir).unwrap(); fs::write(claude_dir.join("settings.json"), r#"{"model": "opus"}"#).unwrap(); - tracevault_cli::commands::init::init_in_directory(tmp.path(), None) + tracevault_cli::commands::init::init_in_directory(tmp.path(), None, None) .await .unwrap(); @@ -99,7 +99,7 @@ fn tracevault_hooks_has_pre_post_and_notification() { async fn init_installs_git_pre_push_hook() { let tmp = tmp_git_repo(); - tracevault_cli::commands::init::init_in_directory(tmp.path(), None) + tracevault_cli::commands::init::init_in_directory(tmp.path(), None, None) .await .unwrap(); @@ -127,7 +127,7 @@ async fn init_preserves_existing_pre_push_hook() { ) .unwrap(); - tracevault_cli::commands::init::init_in_directory(tmp.path(), None) + tracevault_cli::commands::init::init_in_directory(tmp.path(), None, None) .await .unwrap(); @@ -144,10 +144,10 @@ async fn init_preserves_existing_pre_push_hook() { async fn init_does_not_duplicate_hook_on_reinit() { let tmp = tmp_git_repo(); - tracevault_cli::commands::init::init_in_directory(tmp.path(), None) + tracevault_cli::commands::init::init_in_directory(tmp.path(), None, None) .await .unwrap(); - tracevault_cli::commands::init::init_in_directory(tmp.path(), None) + tracevault_cli::commands::init::init_in_directory(tmp.path(), None, None) .await .unwrap(); @@ -163,7 +163,7 @@ async fn init_does_not_duplicate_hook_on_reinit() { async fn init_installs_post_commit_hook() { let tmp = tmp_git_repo(); - tracevault_cli::commands::init::init_in_directory(tmp.path(), None) + tracevault_cli::commands::init::init_in_directory(tmp.path(), None, None) .await .unwrap(); @@ -180,10 +180,10 @@ async fn init_installs_post_commit_hook() { async fn init_does_not_duplicate_post_commit_hook_on_reinit() { let tmp = tmp_git_repo(); - tracevault_cli::commands::init::init_in_directory(tmp.path(), None) + tracevault_cli::commands::init::init_in_directory(tmp.path(), None, None) .await .unwrap(); - tracevault_cli::commands::init::init_in_directory(tmp.path(), None) + tracevault_cli::commands::init::init_in_directory(tmp.path(), None, None) .await .unwrap(); @@ -199,9 +199,13 @@ async fn init_does_not_duplicate_post_commit_hook_on_reinit() { async fn init_writes_server_url_to_config() { let tmp = tmp_git_repo(); - tracevault_cli::commands::init::init_in_directory(tmp.path(), Some("https://tv.example.com")) - .await - .unwrap(); + tracevault_cli::commands::init::init_in_directory( + tmp.path(), + Some("https://tv.example.com"), + None, + ) + .await + .unwrap(); let config_path = tmp.path().join(".tracevault/config.toml"); let content = fs::read_to_string(&config_path).unwrap();