From 8209c2c12363767893f4c4777573cdd3171e6d54 Mon Sep 17 00:00:00 2001 From: winjer Date: Mon, 11 May 2026 09:53:03 +0100 Subject: [PATCH 01/10] feat: add OpenCode source support Adds OpenCode as a second trace source alongside Claude Code. Introduces Source enum, multi-source Input dispatch, opencode message normalizer, pi-ready state keys, and an OpenCode TypeScript plugin. Co-Authored-By: Claude Sonnet 4.6 --- AGENTS.md | 1 + CLAUDE.md | 49 ++++++-- README.md | 85 ++++++++++++-- install.sh | 57 ++++++++- plugin/README.md | 100 ++++++++++++++++ plugin/code-trace.ts | 154 ++++++++++++++++++++++++ src/emit.rs | 24 +++- src/lib.rs | 2 + src/log.rs | 9 +- src/main.rs | 198 ++++++++++++++++++++----------- src/opencode.rs | 240 ++++++++++++++++++++++++++++++++++++++ src/payload.rs | 156 ++++++++++++++++++++++--- src/source.rs | 48 ++++++++ src/state.rs | 59 +++++++--- src/tags.rs | 60 ++++++++-- tests/integration_test.rs | 44 +++++++ 16 files changed, 1147 insertions(+), 139 deletions(-) create mode 120000 AGENTS.md create mode 100644 plugin/README.md create mode 100644 plugin/code-trace.ts create mode 100644 src/opencode.rs create mode 100644 src/source.rs diff --git a/AGENTS.md b/AGENTS.md new file mode 120000 index 0000000..681311e --- /dev/null +++ b/AGENTS.md @@ -0,0 +1 @@ +CLAUDE.md \ No newline at end of file diff --git a/CLAUDE.md b/CLAUDE.md index 680bdb2..359c554 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,11 +1,46 @@ -This project will be a rust re-implementation of the python script in `docs/langfuse_hook.py`. +This project sends Claude Code and OpenCode session traces to Langfuse for observability. -This web page explains the usage: https://langfuse.com/integrations/other/claude-code +## Architecture -The reason to reimplement in rust is: +- Single Rust binary that supports two input sources +- Claude Code: invoked as a Stop hook, reads transcript JSONL +- OpenCode: invoked by a TypeScript plugin, receives messages via stdin -1. I want a really simple way to install this for our teams of developers -2. I want to avoid python environment messing about -3. Each turn takes 0.5 - 1.0s to send the trace, which will be annoying soon. I want to look at ways of speeding this up as much as possible. +## Key files -In addition there is a bug in the current python script where tags are not being successfully sent. It is not clear why but it smells like a bug in the new langfuse python sdk v4. +- `src/main.rs` — entry point, dispatches on `Input` enum +- `src/source.rs` — `Source` enum (ClaudeCode | Opencode) +- `src/payload.rs` — parses stdin into `Input` enum +- `src/opencode.rs` — normalizes OpenCode SDK message format to Claude-format Values +- `src/transcript.rs` — reads Claude Code JSONL transcript +- `src/turns.rs` — groups messages into user/assistant/tool turns +- `src/emit.rs` — builds Langfuse ingestion batch +- `src/tags.rs` — gathers env tags (repo, branch, user, host, os, agent version) +- `src/state.rs` — persisted cursor state per session +- `src/log.rs` — logging to `~/.local/share/code-trace/` + +## State location + +State moved from `~/.claude/state/` to `~/.local/share/code-trace/`. Migration happens on first run. + +## Building + +```bash +cargo build --release +``` + +## Testing + +```bash +cargo test +``` + +## Adding a new source + +1. Add variant to `src/source.rs` `Source` enum +2. Add variant to `src/payload.rs` `Input` enum with parsing +3. Add message normalizer (e.g. `src/opencode.rs`) if needed +4. Update `src/tags.rs` `gather_env_tags` for source-specific tags +5. Update `src/emit.rs` `build_ingestion_batch` for source-specific trace name/metadata +6. Update `src/main.rs` match arm for the new source +7. Update `tests/integration_test.rs` with fixture diff --git a/README.md b/README.md index d77f143..6e749de 100644 --- a/README.md +++ b/README.md @@ -1,11 +1,18 @@ # code-trace -Send [Claude Code](https://docs.anthropic.com/en/docs/claude-code) session traces to [Langfuse](https://langfuse.com) for observability. +Send [Claude Code](https://docs.anthropic.com/en/docs/claude-code) or [OpenCode](https://opencode.ai) session traces to [Langfuse](https://langfuse.com) for observability. -Runs as a Claude Code [Stop hook](https://docs.anthropic.com/en/docs/claude-code/hooks) — after each assistant response, it reads the session transcript, assembles conversational turns, and sends them to Langfuse as structured traces with generations and tool spans. +Runs as a Claude Code [Stop hook](https://docs.anthropic.com/en/docs/claude-code/hooks) or an OpenCode plugin — after each assistant response, it reads the session transcript, assembles conversational turns, and sends them to Langfuse as structured traces with generations and tool spans. Written in Rust for fast startup and zero runtime dependencies. The process forks after assembling the payload — the parent exits immediately while the child sends the HTTP request in the background, adding minimal latency to your workflow. +## Supported agents + +| Agent | Integration | +|-------|-------------| +| Claude Code | Stop hook (settings.json) | +| OpenCode | Plugin (`.opencode/plugins/` or npm) | + ## What you get in Langfuse Each turn produces: @@ -24,6 +31,7 @@ Traces are automatically tagged with: | `host:` | `host:codex` | | `os:` | `os:linux` | | `cc-version:` | `cc-version:2.1.100 (Claude Code)` | +| `oc-version:` | `oc-version:0.4.5 (OpenCode)` | ## Install @@ -35,6 +43,12 @@ curl -sfL https://raw.githubusercontent.com/isotoma/code-trace/main/install.sh | This installs the binary to `~/.local/bin/code-trace`. Make sure `~/.local/bin` is in your `PATH`. +To also install the OpenCode plugin: + +```bash +curl -sfL https://raw.githubusercontent.com/isotoma/code-trace/main/install.sh | bash -s -- --opencode +``` + ### From source ```bash @@ -52,9 +66,11 @@ cp target/release/code-trace ~/.local/bin/ ## Configuration -### 1. Register the hook +### Claude Code -Add to `~/.claude/settings.json`: +#### 1. Register the hook + +The install script does this automatically. If not, add to `~/.claude/settings.json`: ```json { @@ -73,7 +89,7 @@ Add to `~/.claude/settings.json`: } ``` -### 2. Set credentials per project +#### 2. Set credentials per project Add to `.claude/settings.local.json` in your project root: @@ -89,6 +105,36 @@ Add to `.claude/settings.local.json` in your project root: Or set them globally if you want tracing on all projects. +### OpenCode + +#### 1. Install the plugin + +Copy `plugin/code-trace.ts` to your OpenCode plugins directory: + +```bash +mkdir -p ~/.config/opencode/plugins/ +cp plugin/code-trace.ts ~/.config/opencode/plugins/code-trace.ts +``` + +Or add to your `opencode.json`: + +```json +{ + "plugin": ["code-trace"] +} +``` + +#### 2. Set environment variables + +Enable tracing in your shell profile (`.bashrc`, `.zshrc`, etc.): + +```bash +export TRACE_TO_LANGFUSE=true +export LANGFUSE_PUBLIC_KEY=pk-lf-... +export LANGFUSE_SECRET_KEY=sk-lf-... +export LANGFUSE_BASE_URL=https://cloud.langfuse.com # optional, defaults to cloud.langfuse.com +``` + ## Environment variables | Variable | Required | Description | @@ -97,21 +143,42 @@ Or set them globally if you want tracing on all projects. | `LANGFUSE_PUBLIC_KEY` | Yes | Langfuse public key | | `LANGFUSE_SECRET_KEY` | Yes | Langfuse secret key | | `LANGFUSE_BASE_URL` | No | Langfuse host (default: `https://cloud.langfuse.com`) | -| `CC_TRACE_DEBUG` | No | Set to `true` for debug logging | +| `CODE_TRACE_DEBUG` | No | Set to `true` for debug logging (alias: `CC_TRACE_DEBUG`) | The `CC_LANGFUSE_` prefix is also accepted for all Langfuse variables (e.g. `CC_LANGFUSE_PUBLIC_KEY`). ## How it works +### Claude Code + 1. Claude Code calls the hook after each assistant response, passing a JSON payload on stdin with the `sessionId` and `transcriptPath` -2. The binary reads new lines from the transcript JSONL file (tracking offset in `~/.claude/state/code_trace_state.json`) +2. The binary reads new lines from the transcript JSONL file (tracking offset in `~/.local/share/code-trace/state.json`) 3. Messages are grouped into turns (user message + assistant responses + tool results) 4. A batch of Langfuse ingestion events is built (`trace-create`, `generation-create`, `span-create`) 5. The process forks — the parent exits immediately, while the child sends the batch to the Langfuse API via HTTP and logs the result -## Logs +### OpenCode + +1. The OpenCode plugin hooks into the `session.idle` event after each assistant response +2. It fetches new messages since the last processed message (tracked per-session in `~/.local/share/code-trace/opencode_cursor.json`) +3. Messages are assembled into turns and piped to the `code-trace` binary over stdin +4. The binary forks — the parent returns immediately, while the child sends the batch to the Langfuse API via HTTP and logs the result + +## State and logs + +State is stored in `~/.local/share/code-trace/`: +- `state.json` — turn cursor per session +- `state.lock` — file lock for concurrent access +- `opencode_cursor.json` — OpenCode per-session message cursor +- `code_trace.log` — trace log + +Note: State was previously stored in `~/.claude/state/`. On first run, existing state is migrated automatically. + +Enable debug logging with `CODE_TRACE_DEBUG=true`. + +## State migration -Logs are written to `~/.claude/state/code_trace.log`. Enable debug logging with `CC_TRACE_DEBUG=true`. +On first run after an update, any existing state in `~/.claude/state/code_trace_state.json` is migrated to `~/.local/share/code-trace/state.json`. ## License diff --git a/install.sh b/install.sh index 28e0687..d03d709 100755 --- a/install.sh +++ b/install.sh @@ -5,6 +5,17 @@ REPO="isotoma/code-trace" BINARY="code-trace" INSTALL_DIR="${HOME}/.local/bin" SETTINGS_FILE="${HOME}/.claude/settings.json" +OPENCODE_PLUGIN_DIR="${HOME}/.config/opencode/plugins" + +# Parse flags +INSTALL_OPENCODE=false +if [ "${1:-}" = "--opencode" ] || [ "${1:-}" = "-o" ]; then + INSTALL_OPENCODE=true +fi + +detect_opencode() { + [ -d "${HOME}/.config/opencode" ] || [ -f "${HOME}/.config/opencode/opencode.json" ] +} # Detect platform OS="$(uname -s | tr '[:upper:]' '[:lower:]')" @@ -65,15 +76,22 @@ if ! echo "${PATH}" | tr ':' '\n' | grep -qx "${INSTALL_DIR}"; then fi fi -# Register the Claude Code hook +# Determine plugin source location +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PLUGIN_SRC="" +if [ -f "${SCRIPT_DIR}/plugin/code-trace.ts" ]; then + PLUGIN_SRC="${SCRIPT_DIR}/plugin/code-trace.ts" +elif [ -f "${SCRIPT_DIR}/../plugin/code-trace.ts" ]; then + PLUGIN_SRC="$(cd "${SCRIPT_DIR}/../plugin" && pwd)/code-trace.ts" +fi + +# Register Claude Code hook HOOK_ENTRY='{"type":"command","command":"code-trace"}' if [ -f "${SETTINGS_FILE}" ]; then - # Check if code-trace hook is already registered if grep -q "code-trace" "${SETTINGS_FILE}"; then echo "Hook already registered in ${SETTINGS_FILE}" else - # Merge the hook into existing settings using python (available on macOS and most Linux) python3 -c " import json, sys @@ -84,7 +102,6 @@ hook = {'type': 'command', 'command': 'code-trace'} hooks = settings.setdefault('hooks', {}) stop = hooks.setdefault('Stop', []) -# Find or create the hooks list entry for entry in stop: if 'hooks' in entry: entry['hooks'].append(hook) @@ -118,8 +135,36 @@ EOF echo "Created ${SETTINGS_FILE} with hook" fi +# Install OpenCode plugin +install_opencode_plugin() { + if [ -z "${PLUGIN_SRC}" ] || [ ! -f "${PLUGIN_SRC}" ]; then + echo "" + echo "Note: Plugin source not found at ${PLUGIN_SRC}. Skipping OpenCode plugin install." + echo "You can manually copy plugin/code-trace.ts to ${OPENCODE_PLUGIN_DIR}/" + return + fi + + mkdir -p "${OPENCODE_PLUGIN_DIR}" + cp "${PLUGIN_SRC}" "${OPENCODE_PLUGIN_DIR}/code-trace.ts" + echo "Installed OpenCode plugin to ${OPENCODE_PLUGIN_DIR}/code-trace.ts" +} + +if [ "${INSTALL_OPENCODE}" = true ]; then + install_opencode_plugin +elif detect_opencode; then + echo "" + echo "OpenCode detected. Install the code-trace plugin?" + echo " ${OPENCODE_PLUGIN_DIR}/code-trace.ts" + echo "" + read -p "Install OpenCode plugin? [y/N] " -r + if [[ "${REPLY}" =~ ^[Yy]$ ]]; then + install_opencode_plugin + fi +fi + echo "" -echo "Done! To enable tracing, add to your project's .claude/settings.local.json:" +echo "Done! To enable tracing, add to your project's .claude/settings.local.json (Claude Code)" +echo "or set environment variables for OpenCode:" echo "" cat << 'EOF' { @@ -130,3 +175,5 @@ cat << 'EOF' } } EOF +echo "" +echo "For OpenCode, set these environment variables in your shell profile." diff --git a/plugin/README.md b/plugin/README.md new file mode 100644 index 0000000..6db8438 --- /dev/null +++ b/plugin/README.md @@ -0,0 +1,100 @@ +# code-trace OpenCode Plugin + +Sends [OpenCode](https://opencode.ai) session traces to [Langfuse](https://langfuse.com) for observability. + +This is the companion plugin for the [`code-trace`](https://github.com/isotoma/code-trace) Rust binary. + +## Setup + +### 1. Install code-trace binary + +```bash +curl -sfL https://raw.githubusercontent.com/isotoma/code-trace/main/install.sh | bash +``` + +Or build from source: + +```bash +cargo install --git https://github.com/isotoma/code-trace.git +``` + +### 2. Install the plugin + +Copy `code-trace.ts` to your OpenCode plugins directory: + +```bash +mkdir -p ~/.config/opencode/plugins/ +cp code-trace.ts ~/.config/opencode/plugins/code-trace.ts +``` + +Or add to your `opencode.json`: + +```json +{ + "plugin": ["code-trace"] +} +``` + +### 3. Set environment variables + +Enable tracing and set your Langfuse credentials: + +```bash +export TRACE_TO_LANGFUSE=true +export LANGFUSE_PUBLIC_KEY=pk-lf-... +export LANGFUSE_SECRET_KEY=sk-lf-... +export LANGFUSE_BASE_URL=https://cloud.langfuse.com # optional, defaults to cloud.langfuse.com +``` + +Or add them to your shell profile (`.bashrc`, `.zshrc`, etc.). + +## How it works + +The plugin hooks into OpenCode's `session.idle` event, which fires after each assistant response. + +On each idle event: + +1. The plugin fetches new messages since the last processed message (tracked per-session in `~/.local/share/code-trace/opencode_cursor.json`) +2. Messages are assembled into turns (user input + assistant responses + tool results) +3. A JSON payload is sent to the `code-trace` binary over stdin +4. The binary forks — the parent returns immediately while the child sends the trace to Langfuse + +## What you get in Langfuse + +Each turn produces: + +- **Trace** with session grouping, input/output, and tags +- **Generation** span with model name and token-level cost tracking +- **Tool spans** for each tool call (Bash, Read, Edit, etc.) with input/output + +Traces are automatically tagged with: + +| Tag | Example | +|-----|---------| +| `repo:` | `repo:my-project` | +| `branch:` | `branch:main` | +| `user:` | `user:doug` | +| `host:` | `host:codex` | +| `os:` | `os:linux` | +| `oc-version:` | `oc-version:0.4.5` | + +## Troubleshooting + +Enable debug logging: + +```bash +export CODE_TRACE_DEBUG=true +``` + +Then check the log file at `~/.local/share/code-trace/code_trace.log`. + +## Environment Variables + +| Variable | Required | Description | +|----------|----------|-------------| +| `TRACE_TO_LANGFUSE` | Yes | Set to `true` to enable tracing | +| `LANGFUSE_PUBLIC_KEY` | Yes | Langfuse public key | +| `LANGFUSE_SECRET_KEY` | Yes | Langfuse secret key | +| `LANGFUSE_BASE_URL` | No | Langfuse host (default: `https://cloud.langfuse.com`) | +| `CODE_TRACE_DEBUG` | No | Set to `true` for debug logging | +| `CODE_TRACE_BIN` | No | Path to `code-trace` binary (default: `code-trace` on PATH) | diff --git a/plugin/code-trace.ts b/plugin/code-trace.ts new file mode 100644 index 0000000..fe98017 --- /dev/null +++ b/plugin/code-trace.ts @@ -0,0 +1,154 @@ +import { spawn } from "node:child_process"; +import { readFileSync, writeFileSync, existsSync, mkdirSync } from "node:fs"; +import { join, dirname } from "node:path"; +import { fileURLToPath } from "node:url"; + +interface CursorEntry { + lastIndex: number; + lastId: string; +} + +interface CursorStore { + [sessionId: string]: CursorEntry; +} + +function getCursorPath(): string { + const home = process.env.HOME ?? process.env.USERPROFILE ?? "~"; + return join(home, ".local", "share", "code-trace", "opencode_cursor.json"); +} + +function ensureDir(path: string): void { + const dir = dirname(path); + if (!existsSync(dir)) { + mkdirSync(dir, { recursive: true }); + } +} + +function loadCursor(): CursorStore { + const path = getCursorPath(); + try { + if (existsSync(path)) { + return JSON.parse(readFileSync(path, "utf-8")); + } + } catch { + // ignore + } + return {}; +} + +function saveCursor(store: CursorStore): void { + const path = getCursorPath(); + ensureDir(path); + writeFileSync(path, JSON.stringify(store, null, 2)); +} + +function getOpencodeVersion(): string | undefined { + try { + const { execSync } = require("node:child_process"); + const result = execSync("opencode --version", { encoding: "utf-8", timeout: 5000 }); + return result.trim() || undefined; + } catch { + return undefined; + } +} + +export const CodeTracePlugin = async (ctx: { + project?: unknown; + client: { + session: { + messages: (opts: { path: { id: string } }) => Promise<{ + data: Array<{ info: { id: string; role: string; model?: string }; parts: unknown[] }>; + }>; + }; + app: { + log: (opts: { body: { service: string; level: string; message: string; extra?: unknown } }) => Promise; + }; + }; + directory: string; + worktree: string; + $: { + command: (cmd: string, args?: string[]) => { stdout: { text: () => string } }; + }; +}) => { + return { + event: async (event: { type: string; properties?: Record }) => { + if (event.type !== "session.idle") return; + + if (process.env.TRACE_TO_LANGFUSE?.toLowerCase() !== "true") return; + + const sessionId = event.properties?.sessionID as string | undefined; + if (!sessionId) { + await ctx.client.app.log({ + body: { + service: "code-trace", + level: "warn", + message: "session.idle event missing sessionID", + }, + }); + return; + } + + const cursor = loadCursor(); + const prev = cursor[sessionId]; + const startIndex = prev?.lastIndex ?? 0; + + let messagesResponse; + try { + messagesResponse = await ctx.client.session.messages({ path: { id: sessionId } }); + } catch (err) { + await ctx.client.app.log({ + body: { + service: "code-trace", + level: "error", + message: `Failed to fetch session messages: ${err}`, + }, + }); + return; + } + + const allMessages = messagesResponse.data; + if (allMessages.length <= startIndex) return; + + const newMessages = allMessages.slice(startIndex); + if (newMessages.length === 0) return; + + const lastMsg = newMessages[newMessages.length - 1]; + const lastId = lastMsg?.info?.id ?? String(startIndex + newMessages.length - 1); + const agentVersion = getOpencodeVersion(); + + const payload = { + source: "opencode", + sessionId, + cwd: ctx.directory, + messages: newMessages, + agentVersion, + }; + + const binPath = process.env.CODE_TRACE_BIN ?? "code-trace"; + + try { + const child = spawn(binPath, [], { + stdio: ["pipe", "ignore", "ignore"], + detached: true, + shell: true, + }); + + child.stdin?.end(JSON.stringify(payload), "utf-8"); + child.stdin?.destroy(); + child.unref(); + } catch (err) { + await ctx.client.app.log({ + body: { + service: "code-trace", + level: "error", + message: `Failed to spawn code-trace: ${err}`, + }, + }); + return; + } + + cursor[sessionId] = { lastIndex: allMessages.length, lastId }; + saveCursor(cursor); + }, + }; +}; diff --git a/src/emit.rs b/src/emit.rs index 99ff8ac..594b18d 100644 --- a/src/emit.rs +++ b/src/emit.rs @@ -1,4 +1,5 @@ use crate::log; +use crate::source::Source; use crate::transcript; use crate::turns::Turn; use serde_json::{json, Value}; @@ -61,6 +62,7 @@ pub fn build_ingestion_batch( turn: &Turn, transcript_path: &Path, tags: &[String], + source: Source, ) -> Vec { let user_text_raw = transcript::extract_text(transcript::get_content(&turn.user_msg)); let (user_text, user_text_meta) = truncate(&user_text_raw); @@ -85,13 +87,13 @@ pub fn build_ingestion_batch( "body": { "id": trace_id, "timestamp": now, - "name": format!("Claude Code - Turn {turn_num}"), + "name": format!("{} - Turn {turn_num}", source.trace_name_prefix()), "sessionId": session_id, "input": json!({"role": "user", "content": user_text}), "output": json!({"role": "assistant", "content": assistant_text}), "tags": tags, "metadata": { - "source": "claude-code", + "source": source.as_str(), "session_id": session_id, "turn_number": turn_num, "transcript_path": transcript_path.to_string_lossy(), @@ -264,8 +266,9 @@ mod tests { #[test] fn builds_trace_and_generation_events() { + use crate::source::Source; let turn = make_simple_turn(); - let events = build_ingestion_batch("sess1", 1, &turn, Path::new("/tmp/t.jsonl"), &["claude-code".to_string()]); + let events = build_ingestion_batch("sess1", 1, &turn, Path::new("/tmp/t.jsonl"), &["claude-code".to_string()], Source::ClaudeCode); assert_eq!(events.len(), 2); assert_eq!(events[0]["type"], "trace-create"); assert_eq!(events[1]["type"], "generation-create"); @@ -273,9 +276,10 @@ mod tests { #[test] fn trace_event_has_tags() { + use crate::source::Source; let turn = make_simple_turn(); let tags = vec!["claude-code".to_string(), "repo:myrepo".to_string()]; - let events = build_ingestion_batch("sess1", 1, &turn, Path::new("/tmp/t.jsonl"), &tags); + let events = build_ingestion_batch("sess1", 1, &turn, Path::new("/tmp/t.jsonl"), &tags, Source::ClaudeCode); let trace_tags = events[0]["body"]["tags"].as_array().unwrap(); assert_eq!(trace_tags.len(), 2); assert_eq!(trace_tags[0], "claude-code"); @@ -284,6 +288,7 @@ mod tests { #[test] fn builds_tool_spans() { + use crate::source::Source; let mut tool_results = HashMap::new(); tool_results.insert("tu_1".to_string(), json!("file1.txt\nfile2.txt")); let turn = Turn { @@ -294,12 +299,21 @@ mod tests { ]}})], tool_results_by_id: tool_results, }; - let events = build_ingestion_batch("sess1", 1, &turn, Path::new("/tmp/t.jsonl"), &["claude-code".to_string()]); + let events = build_ingestion_batch("sess1", 1, &turn, Path::new("/tmp/t.jsonl"), &["claude-code".to_string()], Source::ClaudeCode); assert_eq!(events.len(), 3); assert_eq!(events[2]["type"], "span-create"); assert_eq!(events[2]["body"]["name"], "Tool: Bash"); } + #[test] + fn opencode_source_uses_correct_trace_name() { + use crate::source::Source; + let turn = make_simple_turn(); + let events = build_ingestion_batch("sess1", 1, &turn, Path::new("/tmp/t.jsonl"), &["opencode".to_string()], Source::Opencode); + assert_eq!(events[0]["body"]["name"], "OpenCode - Turn 1"); + assert_eq!(events[0]["body"]["metadata"]["source"], "opencode"); + } + #[test] fn truncates_long_text() { let long = "x".repeat(30_000); diff --git a/src/lib.rs b/src/lib.rs index 25b1027..709e4e7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,8 @@ pub mod emit; pub mod log; +pub mod opencode; pub mod payload; +pub mod source; pub mod state; pub mod tags; pub mod transcript; diff --git a/src/log.rs b/src/log.rs index 373986c..5ce2105 100644 --- a/src/log.rs +++ b/src/log.rs @@ -5,7 +5,7 @@ use std::path::PathBuf; const MAX_LOG_BYTES: u64 = 512 * 1024; // 512KB fn log_dir() -> Option { - let dir = dirs::home_dir()?.join(".claude").join("state"); + let dir = dirs::data_local_dir()?.join("code-trace"); fs::create_dir_all(&dir).ok()?; Some(dir) } @@ -42,7 +42,12 @@ pub fn error(msg: &str) { } pub fn debug(msg: &str) { - if std::env::var("CC_TRACE_DEBUG").unwrap_or_default().to_lowercase() == "true" { + let debug_enabled = std::env::var("CODE_TRACE_DEBUG") + .or_else(|_| std::env::var("CC_TRACE_DEBUG")) + .unwrap_or_default() + .to_lowercase() + == "true"; + if debug_enabled { write_log("DEBUG", msg); } } diff --git a/src/main.rs b/src/main.rs index c17497d..4d9b1e4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,11 +1,10 @@ -use code_trace::{emit, log, payload, state, tags, transcript, turns}; +use code_trace::{emit, log, opencode, payload, state, tags, transcript, turns}; use std::time::Instant; fn run() -> i32 { let start = Instant::now(); log::debug("code-trace started"); - // Check enable flag if std::env::var("TRACE_TO_LANGFUSE") .unwrap_or_default() .to_lowercase() @@ -14,7 +13,6 @@ fn run() -> i32 { return 0; } - // Read config from env let public_key = std::env::var("CC_LANGFUSE_PUBLIC_KEY") .or_else(|_| std::env::var("LANGFUSE_PUBLIC_KEY")) .unwrap_or_default(); @@ -35,79 +33,147 @@ fn run() -> i32 { secret_key, }; - // Read hook payload from stdin let raw = payload::read_stdin(); - let hook = payload::parse_payload(&raw); - - let Some(session_id) = hook.session_id else { - log::debug("Missing session_id; exiting"); - return 0; - }; - let Some(transcript_path) = hook.transcript_path else { - log::debug("Missing transcript_path; exiting"); - return 0; + let input = payload::parse_payload(&raw); + + let source = input.source(); + let session_id = match input.session_id() { + Some(s) => s.to_string(), + None => { + log::debug("Missing session_id; exiting"); + return 0; + } }; - if !transcript_path.exists() { - log::debug(&format!( - "Transcript does not exist: {}", - transcript_path.display() - )); - return 0; - } - - // Gather tags - let env_tags = tags::gather_env_tags(hook.cwd.as_deref()); + let cwd = input.cwd().map(String::from); + let env_tags = tags::gather_env_tags(source, cwd.as_deref(), input.agent_version()); - // Lock, load state, read new messages, build turns, emit let _lock = state::FileLock::acquire(); let mut global_state = state::load_state(); state::prune_old_sessions(&mut global_state); - let key = state::state_key(&session_id, &transcript_path.to_string_lossy()); - let mut ss = global_state - .get(&key) - .cloned() - .unwrap_or_default(); - let msgs = transcript::read_new_jsonl(&transcript_path, &mut ss); - if msgs.is_empty() { - global_state.insert(key, ss); - state::save_state(&global_state); - return 0; + match input { + payload::Input::ClaudeCode { + session_id: _, + transcript_path, + cwd: _, + } => { + let Some(transcript_path) = transcript_path else { + log::debug("Missing transcript_path; exiting"); + return 0; + }; + + if !transcript_path.exists() { + log::debug(&format!( + "Transcript does not exist: {}", + transcript_path.display() + )); + return 0; + } + + let key = state::state_key( + source.as_str(), + &session_id, + &transcript_path.to_string_lossy(), + ); + let mut ss = global_state.get(&key).cloned().unwrap_or_default(); + + let msgs = transcript::read_new_jsonl(&transcript_path, &mut ss); + if msgs.is_empty() { + global_state.insert(key, ss); + state::save_state(&global_state); + return 0; + } + + let built_turns = turns::build_turns(msgs); + if built_turns.is_empty() { + global_state.insert(key, ss); + state::save_state(&global_state); + return 0; + } + + let mut all_events = Vec::new(); + let mut emitted = 0u32; + for t in &built_turns { + emitted += 1; + let turn_num = ss.turn_count + emitted; + let events = emit::build_ingestion_batch( + &session_id, + turn_num, + t, + &transcript_path, + &env_tags, + source, + ); + all_events.extend(events); + } + + ss.turn_count += emitted; + state::touch(&mut ss); + global_state.insert(key, ss); + state::save_state(&global_state); + + emit::send_batch_fire_and_forget(&config, all_events); + + let dur = start.elapsed(); + log::info(&format!( + "Processed {emitted} turns in {:.2}s (session={session_id})", + dur.as_secs_f64() + )); + } + + payload::Input::Opencode { + session_id: _, + cwd: _, + messages, + agent_version: _, + } => { + if messages.is_empty() { + return 0; + } + + let key = state::state_key(source.as_str(), &session_id, &session_id); + let mut ss = global_state.get(&key).cloned().unwrap_or_default(); + + let normalized = opencode::normalize_opencode_messages(messages); + let built_turns = turns::build_turns(normalized); + if built_turns.is_empty() { + global_state.insert(key, ss); + state::save_state(&global_state); + return 0; + } + + let mut all_events = Vec::new(); + let mut emitted = 0u32; + for t in &built_turns { + emitted += 1; + let turn_num = ss.turn_count + emitted; + let events = emit::build_ingestion_batch( + &session_id, + turn_num, + t, + std::path::Path::new("opencode"), + &env_tags, + source, + ); + all_events.extend(events); + } + + ss.turn_count += emitted; + state::touch(&mut ss); + global_state.insert(key, ss); + state::save_state(&global_state); + + emit::send_batch_fire_and_forget(&config, all_events); + + let dur = start.elapsed(); + log::info(&format!( + "Processed {emitted} turns in {:.2}s (session={session_id})", + dur.as_secs_f64() + )); + } } - let built_turns = turns::build_turns(msgs); - if built_turns.is_empty() { - global_state.insert(key, ss); - state::save_state(&global_state); - return 0; - } - - // Collect all events across turns into one batch - let mut all_events = Vec::new(); - let mut emitted = 0u32; - for t in &built_turns { - emitted += 1; - let turn_num = ss.turn_count + emitted; - let events = - emit::build_ingestion_batch(&session_id, turn_num, t, &transcript_path, &env_tags); - all_events.extend(events); - } - - ss.turn_count += emitted; - state::touch(&mut ss); - global_state.insert(key, ss); - state::save_state(&global_state); - - // Fire and forget - emit::send_batch_fire_and_forget(&config, all_events); - - let dur = start.elapsed(); - log::info(&format!( - "Processed {emitted} turns in {:.2}s (session={session_id})", - dur.as_secs_f64() - )); - 0 } diff --git a/src/opencode.rs b/src/opencode.rs new file mode 100644 index 0000000..77367a9 --- /dev/null +++ b/src/opencode.rs @@ -0,0 +1,240 @@ +use serde_json::{json, Value}; +use std::collections::HashMap; + +pub fn normalize_opencode_messages(messages: Vec) -> Vec { + let mut result: Vec = Vec::new(); + let mut pending_tool_results: HashMap = HashMap::new(); + let mut pending_assistant_idx: Option = None; + + for msg in messages { + let role = get_opencode_role(&msg); + let parts = get_opencode_parts(&msg); + + if role == Some("user") { + let mut user_content: Vec = Vec::new(); + + for part in &parts { + if part.get("type").and_then(|v| v.as_str()) == Some("tool_result") { + let tool_use_id = part + .get("tool_use_id") + .and_then(|v| v.as_str()) + .unwrap_or(""); + let content = part.get("content").cloned().unwrap_or(json!(null)); + pending_tool_results.insert(tool_use_id.to_string(), content); + } + } + + if !pending_tool_results.is_empty() { + let mut tr_arr: Vec = Vec::new(); + for (tid, content) in pending_tool_results.drain() { + tr_arr.push(json!({ + "type": "tool_result", + "tool_use_id": tid, + "content": content, + })); + } + user_content.extend(tr_arr); + } + + let text_parts: Vec = parts + .iter() + .filter(|p| p.get("type").and_then(|v| v.as_str()) == Some("text")) + .map(|p| { + let text = p.get("text").and_then(|v| v.as_str()).unwrap_or(""); + json!({ "type": "text", "text": text }) + }) + .collect(); + if !text_parts.is_empty() { + user_content.extend(text_parts); + } + + if !user_content.is_empty() { + result.push(json!({ + "type": "user", + "message": { + "role": "user", + "content": user_content, + } + })); + } + pending_assistant_idx = None; + continue; + } + + if role == Some("assistant") { + let id = msg + .get("info") + .and_then(|v| v.get("id")) + .and_then(|v| v.as_str()) + .map(String::from) + .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + + let model = msg + .get("info") + .and_then(|v| v.get("model")) + .and_then(|v| v.as_str()) + .unwrap_or("opencode"); + + let mut content_parts: Vec = Vec::new(); + let mut has_tool_use = false; + let mut current_msg_tool_uses: Vec = Vec::new(); + + for part in &parts { + let part_type = part.get("type").and_then(|v| v.as_str()); + match part_type { + Some("text") => { + let text = part.get("text").and_then(|v| v.as_str()).unwrap_or(""); + content_parts.push(json!({ "type": "text", "text": text })); + } + Some("tool_use") => { + has_tool_use = true; + let tool_id = part.get("id").and_then(|v| v.as_str()).unwrap_or(""); + current_msg_tool_uses.push(tool_id.to_string()); + let tool_name = part.get("name").and_then(|v| v.as_str()).unwrap_or("unknown"); + let tool_input = part.get("input").cloned().unwrap_or(json!({})); + content_parts.push(json!({ + "type": "tool_use", + "id": tool_id, + "name": tool_name, + "input": tool_input, + })); + } + Some("tool_result") => { + let tool_use_id = part + .get("tool_use_id") + .and_then(|v| v.as_str()) + .unwrap_or(""); + let content = part.get("content").cloned().unwrap_or(json!(null)); + pending_tool_results.insert(tool_use_id.to_string(), content); + } + _ => {} + } + } + + let assistant_msg_idx = result.len(); + if has_tool_use { + pending_assistant_idx = Some(assistant_msg_idx); + } + + let assistant_msg = json!({ + "type": "assistant", + "message": { + "id": id, + "role": "assistant", + "model": model, + "content": content_parts, + } + }); + + if !pending_tool_results.is_empty() { + if let Some(prev_idx) = pending_assistant_idx { + if prev_idx < result.len() { + let prev_has_tool_use = result[prev_idx]["message"]["content"] + .as_array() + .map(|arr| { + arr.iter().any(|p| { + p.get("type").and_then(|v| v.as_str()) == Some("tool_use") + }) + }) + .unwrap_or(false); + if prev_has_tool_use { + let mut tr_arr: Vec = Vec::new(); + for (tid, content) in pending_tool_results.drain() { + tr_arr.push(json!({ + "type": "tool_result", + "tool_use_id": tid, + "content": content, + })); + } + result[prev_idx]["message"]["tool_results"] = json!(tr_arr); + } + } + } + } + + result.push(assistant_msg); + + if !has_tool_use && !pending_tool_results.is_empty() { + pending_assistant_idx = Some(assistant_msg_idx); + } + continue; + } + + result.push(msg); + } + + result +} + +fn get_opencode_role(msg: &Value) -> Option<&str> { + msg.get("info") + .and_then(|v| v.get("role")) + .and_then(|v| v.as_str()) + .filter(|r| *r == "user" || *r == "assistant") +} + +fn get_opencode_parts(msg: &Value) -> Vec { + msg.get("parts") + .and_then(|v| v.as_array()) + .cloned() + .unwrap_or_default() +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn normalizes_user_message_with_text() { + let msgs = vec![json!({ + "info": { "id": "msg1", "role": "user" }, + "parts": [{ "type": "text", "text": "Hello" }] + })]; + let normalized = normalize_opencode_messages(msgs); + assert_eq!(normalized.len(), 1); + assert_eq!(normalized[0]["type"], "user"); + assert_eq!(normalized[0]["message"]["role"], "user"); + } + + #[test] + fn normalizes_assistant_message_with_tool_use() { + let msgs = vec![json!({ + "info": { "id": "msg2", "role": "assistant", "model": "claude" }, + "parts": [ + { "type": "text", "text": "Let me check" }, + { "type": "tool_use", "id": "tu1", "name": "Bash", "input": { "command": "ls" } } + ] + })]; + let normalized = normalize_opencode_messages(msgs); + assert_eq!(normalized.len(), 1); + assert_eq!(normalized[0]["type"], "assistant"); + assert_eq!(normalized[0]["message"]["content"][1]["type"], "tool_use"); + assert_eq!(normalized[0]["message"]["content"][1]["name"], "Bash"); + } + + #[test] + fn pending_tool_results_attached_to_previous_assistant() { + let msgs = vec![ + json!({ + "info": { "id": "msg1", "role": "assistant" }, + "parts": [ + { "type": "text", "text": "Running command" }, + { "type": "tool_use", "id": "tu1", "name": "Bash", "input": {} } + ] + }), + json!({ + "info": { "id": "msg2", "role": "assistant" }, + "parts": [ + { "type": "tool_result", "tool_use_id": "tu1", "content": "file1.txt" } + ] + }), + ]; + let normalized = normalize_opencode_messages(msgs); + assert_eq!(normalized.len(), 2); + let tool_results = normalized[0]["message"]["tool_results"].as_array().unwrap(); + assert_eq!(tool_results.len(), 1); + assert_eq!(tool_results[0]["tool_use_id"], "tu1"); + assert_eq!(tool_results[0]["content"], "file1.txt"); + } +} diff --git a/src/payload.rs b/src/payload.rs index 5765025..edb587d 100644 --- a/src/payload.rs +++ b/src/payload.rs @@ -1,14 +1,75 @@ use serde_json::Value; use std::path::PathBuf; +use crate::source::Source; + pub struct HookPayload { pub session_id: Option, pub transcript_path: Option, pub cwd: Option, } -/// Parse hook payload from a JSON value. Tolerates multiple field name conventions. -pub fn parse_payload(value: &Value) -> HookPayload { +pub enum Input { + ClaudeCode { + session_id: Option, + transcript_path: Option, + cwd: Option, + }, + Opencode { + session_id: Option, + cwd: Option, + messages: Vec, + agent_version: Option, + }, +} + +impl Input { + pub fn source(&self) -> Source { + match self { + Input::ClaudeCode { .. } => Source::ClaudeCode, + Input::Opencode { .. } => Source::Opencode, + } + } + + pub fn session_id(&self) -> Option<&str> { + match self { + Input::ClaudeCode { session_id, .. } => session_id.as_deref(), + Input::Opencode { session_id, .. } => session_id.as_deref(), + } + } + + pub fn cwd(&self) -> Option<&str> { + match self { + Input::ClaudeCode { cwd, .. } => cwd.as_deref(), + Input::Opencode { cwd, .. } => cwd.as_deref(), + } + } + + pub fn agent_version(&self) -> Option<&str> { + match self { + Input::ClaudeCode { .. } => None, + Input::Opencode { agent_version, .. } => agent_version.as_deref(), + } + } +} + +pub fn parse_payload(value: &Value) -> Input { + if let Some(source_val) = value.get("source").and_then(|v| v.as_str()) { + if let Some(source) = Source::parse(source_val) { + return parse_by_source(&source, value); + } + } + parse_claude_code_payload(value) +} + +fn parse_by_source(source: &Source, value: &Value) -> Input { + match source { + Source::ClaudeCode => parse_claude_code_payload(value), + Source::Opencode => parse_opencode_payload(value), + } +} + +fn parse_claude_code_payload(value: &Value) -> Input { let session_id = value .get("sessionId") .or_else(|| value.get("session_id")) @@ -29,19 +90,44 @@ pub fn parse_payload(value: &Value) -> HookPayload { p }); - let cwd = value - .get("cwd") + let cwd = value.get("cwd").and_then(|v| v.as_str()).map(String::from); + + Input::ClaudeCode { + session_id, + transcript_path: transcript, + cwd, + } +} + +fn parse_opencode_payload(value: &Value) -> Input { + let session_id = value + .get("sessionId") + .or_else(|| value.get("session_id")) + .and_then(|v| v.as_str()) + .map(String::from); + + let cwd = value.get("cwd").and_then(|v| v.as_str()).map(String::from); + + let messages = value + .get("messages") + .and_then(|v| v.as_array()) + .cloned() + .unwrap_or_default(); + + let agent_version = value + .get("agentVersion") + .or_else(|| value.get("agent_version")) .and_then(|v| v.as_str()) .map(String::from); - HookPayload { + Input::Opencode { session_id, - transcript_path: transcript, cwd, + messages, + agent_version, } } -/// Read stdin fully and parse as JSON. Returns Value::Null on any failure. pub fn read_stdin() -> Value { let mut buf = String::new(); if std::io::Read::read_to_string(&mut std::io::stdin(), &mut buf).is_err() { @@ -62,10 +148,9 @@ mod tests { "transcriptPath": "/tmp/test.jsonl", "cwd": "/home/user/project" }); - let p = parse_payload(&v); - assert_eq!(p.session_id.as_deref(), Some("abc-123")); - assert_eq!(p.transcript_path.as_deref(), Some(std::path::Path::new("/tmp/test.jsonl"))); - assert_eq!(p.cwd.as_deref(), Some("/home/user/project")); + let input = parse_payload(&v); + assert_eq!(input.source(), Source::ClaudeCode); + assert_eq!(input.session_id(), Some("abc-123")); } #[test] @@ -74,15 +159,52 @@ mod tests { "session_id": "abc-123", "transcript_path": "/tmp/test.jsonl" }); - let p = parse_payload(&v); - assert_eq!(p.session_id.as_deref(), Some("abc-123")); - assert_eq!(p.transcript_path.as_deref(), Some(std::path::Path::new("/tmp/test.jsonl"))); + let input = parse_payload(&v); + assert_eq!(input.source(), Source::ClaudeCode); + assert_eq!(input.session_id(), Some("abc-123")); + } + + #[test] + fn parses_explicit_claude_code_source() { + let v = json!({ + "source": "claude-code", + "sessionId": "abc-123", + "transcriptPath": "/tmp/test.jsonl" + }); + let input = parse_payload(&v); + assert_eq!(input.source(), Source::ClaudeCode); + } + + #[test] + fn parses_opencode_payload() { + let v = json!({ + "source": "opencode", + "sessionId": "ses_123", + "cwd": "/home/user/project", + "messages": [{"type": "user", "content": "hello"}], + "agentVersion": "0.4.5" + }); + let input = parse_payload(&v); + assert_eq!(input.source(), Source::Opencode); + assert_eq!(input.session_id(), Some("ses_123")); + assert_eq!(input.agent_version(), Some("0.4.5")); + match &input { + Input::Opencode { messages, .. } => assert_eq!(messages.len(), 1), + _ => panic!("expected Opencode input"), + } + } + + #[test] + fn defaults_to_claude_code_when_source_missing() { + let v = json!({"sessionId": "abc"}); + let input = parse_payload(&v); + assert_eq!(input.source(), Source::ClaudeCode); } #[test] fn handles_empty_payload() { - let p = parse_payload(&Value::Null); - assert!(p.session_id.is_none()); - assert!(p.transcript_path.is_none()); + let input = parse_payload(&Value::Null); + assert_eq!(input.source(), Source::ClaudeCode); + assert_eq!(input.session_id(), None); } } diff --git a/src/source.rs b/src/source.rs new file mode 100644 index 0000000..eb4d41b --- /dev/null +++ b/src/source.rs @@ -0,0 +1,48 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] +pub enum Source { + #[default] + ClaudeCode, + Opencode, +} + +impl Source { + pub fn agent_tag(&self) -> &'static str { + match self { + Source::ClaudeCode => "claude-code", + Source::Opencode => "opencode", + } + } + + pub fn version_tag_prefix(&self) -> &'static str { + match self { + Source::ClaudeCode => "cc-version", + Source::Opencode => "oc-version", + } + } + + pub fn trace_name_prefix(&self) -> &'static str { + match self { + Source::ClaudeCode => "Claude Code", + Source::Opencode => "OpenCode", + } + } + + pub fn as_str(&self) -> &'static str { + match self { + Source::ClaudeCode => "claude-code", + Source::Opencode => "opencode", + } + } + + pub fn parse(s: &str) -> Option { + match s.to_lowercase().as_str() { + "claude-code" | "claude_code" | "claude" => Some(Source::ClaudeCode), + "opencode" | "open_code" => Some(Source::Opencode), + _ => None, + } + } +} + + diff --git a/src/state.rs b/src/state.rs index f71ea22..2d63116 100644 --- a/src/state.rs +++ b/src/state.rs @@ -8,17 +8,21 @@ use std::path::PathBuf; use crate::log; fn state_dir() -> Option { - let dir = dirs::home_dir()?.join(".claude").join("state"); + let dir = dirs::data_local_dir()?.join("code-trace"); fs::create_dir_all(&dir).ok()?; Some(dir) } +fn old_state_file() -> Option { + dirs::home_dir().map(|h| h.join(".claude").join("state").join("code_trace_state.json")) +} + fn state_file() -> Option { - Some(state_dir()?.join("code_trace_state.json")) + Some(state_dir()?.join("state.json")) } fn lock_file() -> Option { - Some(state_dir()?.join("code_trace_state.lock")) + Some(state_dir()?.join("state.lock")) } /// RAII file lock using flock. Best-effort: proceeds without lock on failure. @@ -91,9 +95,9 @@ pub fn touch(ss: &mut SessionState) { .as_secs(); } -pub fn state_key(session_id: &str, transcript_path: &str) -> String { +pub fn state_key(source: &str, session_id: &str, handle: &str) -> String { let mut hasher = Sha256::new(); - hasher.update(format!("{session_id}::{transcript_path}")); + hasher.update(format!("{source}:{session_id}:{handle}")); format!("{:x}", hasher.finalize()) } @@ -101,14 +105,30 @@ pub fn load_state() -> GlobalState { let Some(path) = state_file() else { return GlobalState::new(); }; - let Ok(mut file) = File::open(&path) else { - return GlobalState::new(); - }; - let mut buf = String::new(); - if file.read_to_string(&mut buf).is_err() { - return GlobalState::new(); + if path.exists() { + let Ok(mut file) = File::open(&path) else { + return GlobalState::new(); + }; + let mut buf = String::new(); + if file.read_to_string(&mut buf).is_err() { + return GlobalState::new(); + } + return serde_json::from_str(&buf).unwrap_or_default(); } - serde_json::from_str(&buf).unwrap_or_default() + if let Some(old_path) = old_state_file() { + if old_path.exists() { + if let Ok(mut file) = File::open(&old_path) { + let mut buf = String::new(); + if file.read_to_string(&mut buf).is_ok() { + if let Ok(state) = serde_json::from_str::(&buf) { + save_state(&state); + return state; + } + } + } + } + } + GlobalState::new() } pub fn save_state(state: &GlobalState) { @@ -130,16 +150,23 @@ mod tests { #[test] fn state_key_is_deterministic() { - let k1 = state_key("sess1", "/tmp/t.jsonl"); - let k2 = state_key("sess1", "/tmp/t.jsonl"); + let k1 = state_key("claude-code", "sess1", "/tmp/t.jsonl"); + let k2 = state_key("claude-code", "sess1", "/tmp/t.jsonl"); assert_eq!(k1, k2); assert_eq!(k1.len(), 64); // sha256 hex } #[test] fn state_key_differs_for_different_inputs() { - let k1 = state_key("sess1", "/tmp/a.jsonl"); - let k2 = state_key("sess1", "/tmp/b.jsonl"); + let k1 = state_key("claude-code", "sess1", "/tmp/a.jsonl"); + let k2 = state_key("claude-code", "sess1", "/tmp/b.jsonl"); + assert_ne!(k1, k2); + } + + #[test] + fn state_key_differs_for_different_sources() { + let k1 = state_key("claude-code", "sess1", "/tmp/t.jsonl"); + let k2 = state_key("opencode", "sess1", "/tmp/t.jsonl"); assert_ne!(k1, k2); } diff --git a/src/tags.rs b/src/tags.rs index e997ca9..7e613bd 100644 --- a/src/tags.rs +++ b/src/tags.rs @@ -1,5 +1,7 @@ use std::process::Command; +use crate::source::Source; + fn git_cmd(args: &[&str], cwd: Option<&str>) -> Option { let mut cmd = Command::new("git"); cmd.args(args); @@ -19,8 +21,14 @@ fn git_cmd(args: &[&str], cwd: Option<&str>) -> Option { } } -pub fn gather_env_tags(cwd: Option<&str>) -> Vec { - let mut tags = vec!["claude-code".to_string()]; +pub fn gather_env_tags(source: Source, cwd: Option<&str>, agent_version: Option<&str>) -> Vec { + let mut tags = vec![source.agent_tag().to_string()]; + + if let Some(ver) = agent_version { + if !ver.is_empty() { + tags.push(format!("{}:{ver}", source.version_tag_prefix())); + } + } // Git repo name if let Some(toplevel) = git_cmd(&["rev-parse", "--show-toplevel"], cwd) { @@ -47,12 +55,26 @@ pub fn gather_env_tags(cwd: Option<&str>) -> Vec { // OS tags.push(format!("os:{}", std::env::consts::OS)); - // Claude Code version - if let Ok(output) = Command::new("claude").arg("--version").output() { - if output.status.success() { - let ver = String::from_utf8_lossy(&output.stdout).trim().to_string(); - if !ver.is_empty() { - tags.push(format!("cc-version:{ver}")); + // Claude Code version probe (only for Claude Code source, if no version provided) + if source == Source::ClaudeCode && agent_version.is_none() { + if let Ok(output) = Command::new("claude").arg("--version").output() { + if output.status.success() { + let ver = String::from_utf8_lossy(&output.stdout).trim().to_string(); + if !ver.is_empty() { + tags.push(format!("cc-version:{ver}")); + } + } + } + } + + // OpenCode version probe (only for OpenCode source, if no version provided) + if source == Source::Opencode && agent_version.is_none() { + if let Ok(output) = Command::new("opencode").arg("--version").output() { + if output.status.success() { + let ver = String::from_utf8_lossy(&output.stdout).trim().to_string(); + if !ver.is_empty() { + tags.push(format!("oc-version:{ver}")); + } } } } @@ -65,20 +87,34 @@ mod tests { use super::*; #[test] - fn always_includes_claude_code_tag() { - let tags = gather_env_tags(None); + fn always_includes_agent_tag() { + let tags = gather_env_tags(Source::ClaudeCode, None, None); assert!(tags.contains(&"claude-code".to_string())); } #[test] fn includes_os_tag() { - let tags = gather_env_tags(None); + let tags = gather_env_tags(Source::ClaudeCode, None, None); assert!(tags.iter().any(|t| t.starts_with("os:"))); } #[test] fn includes_user_tag() { - let tags = gather_env_tags(None); + let tags = gather_env_tags(Source::ClaudeCode, None, None); assert!(tags.iter().any(|t| t.starts_with("user:"))); } + + #[test] + fn opencode_source_uses_opencode_tag() { + let tags = gather_env_tags(Source::Opencode, None, Some("0.4.5")); + assert!(tags.contains(&"opencode".to_string())); + assert!(tags.contains(&"oc-version:0.4.5".to_string())); + assert!(!tags.iter().any(|t| t.starts_with("cc-version:"))); + } + + #[test] + fn version_from_payload_preferred_over_probe() { + let tags = gather_env_tags(Source::Opencode, None, Some("1.2.3")); + assert!(tags.contains(&"oc-version:1.2.3".to_string())); + } } diff --git a/tests/integration_test.rs b/tests/integration_test.rs index 7bd835b..39b0b78 100644 --- a/tests/integration_test.rs +++ b/tests/integration_test.rs @@ -22,6 +22,7 @@ fn end_to_end_simple_transcript() { &turns[0], transcript.path(), &tags, + code_trace::source::Source::ClaudeCode, ); assert_eq!(events.len(), 2); @@ -52,6 +53,7 @@ fn end_to_end_tool_transcript() { &turns[0], transcript.path(), &["claude-code".to_string()], + code_trace::source::Source::ClaudeCode, ); assert_eq!(events.len(), 3); @@ -59,3 +61,45 @@ fn end_to_end_tool_transcript() { assert_eq!(events[2]["body"]["name"], "Tool: Bash"); assert_eq!(events[2]["body"]["output"], json!("README.md")); } + +#[test] +fn end_to_end_opencode_transcript() { + let msgs = vec![ + json!({ + "info": { "id": "msg_1", "role": "user" }, + "parts": [{ "type": "text", "text": "Hello" }] + }), + json!({ + "info": { "id": "msg_2", "role": "assistant", "model": "claude-sonnet-4-20250514" }, + "parts": [ + { "type": "text", "text": "Hi there!" }, + { "type": "tool_use", "id": "tu_1", "name": "Bash", "input": { "command": "ls" } } + ] + }), + json!({ + "info": { "id": "msg_3", "role": "assistant" }, + "parts": [ + { "type": "tool_result", "tool_use_id": "tu_1", "content": "file1.txt" } + ] + }), + ]; + + let normalized = code_trace::opencode::normalize_opencode_messages(msgs); + let turns = code_trace::turns::build_turns(normalized); + assert_eq!(turns.len(), 1); + + let tags = vec!["opencode".to_string()]; + let events = code_trace::emit::build_ingestion_batch( + "oc-session", + 1, + &turns[0], + std::path::Path::new("opencode"), + &tags, + code_trace::source::Source::Opencode, + ); + + assert_eq!(events.len(), 3); + assert_eq!(events[0]["type"], "trace-create"); + assert_eq!(events[0]["body"]["name"], "OpenCode - Turn 1"); + assert_eq!(events[0]["body"]["metadata"]["source"], "opencode"); +} From 3cee93e9d3dc3cfa2c8faaddcb3c98dd51c3292f Mon Sep 17 00:00:00 2001 From: winjer Date: Mon, 11 May 2026 09:46:39 +0100 Subject: [PATCH 02/10] feat: add PiAgent variant to Source enum --- src/source.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/source.rs b/src/source.rs index eb4d41b..99cf316 100644 --- a/src/source.rs +++ b/src/source.rs @@ -5,6 +5,7 @@ pub enum Source { #[default] ClaudeCode, Opencode, + PiAgent, } impl Source { @@ -12,6 +13,7 @@ impl Source { match self { Source::ClaudeCode => "claude-code", Source::Opencode => "opencode", + Source::PiAgent => "pi-agent", } } @@ -19,6 +21,7 @@ impl Source { match self { Source::ClaudeCode => "cc-version", Source::Opencode => "oc-version", + Source::PiAgent => "pi-version", } } @@ -26,6 +29,7 @@ impl Source { match self { Source::ClaudeCode => "Claude Code", Source::Opencode => "OpenCode", + Source::PiAgent => "Pi Agent", } } @@ -33,6 +37,7 @@ impl Source { match self { Source::ClaudeCode => "claude-code", Source::Opencode => "opencode", + Source::PiAgent => "pi-agent", } } @@ -40,9 +45,8 @@ impl Source { match s.to_lowercase().as_str() { "claude-code" | "claude_code" | "claude" => Some(Source::ClaudeCode), "opencode" | "open_code" => Some(Source::Opencode), + "pi-agent" | "pi_agent" | "pi" => Some(Source::PiAgent), _ => None, } } } - - From cac3aa46cdfe2de00f17942cd3d9f4f30c48517c Mon Sep 17 00:00:00 2001 From: winjer Date: Mon, 11 May 2026 09:48:32 +0100 Subject: [PATCH 03/10] feat: add pi-agent version probe to tags Update gather_env_tags to accept source and agent_version parameters. Add version probing for all three sources (ClaudeCode, Opencode, PiAgent): - ProbeCommand will run only when agent_version is not provided - Push version tag in format "{source.version_tag_prefix()}:{version}" - Include new tests for each source variant Co-Authored-By: Claude Sonnet 4.5 --- src/lib.rs | 1 + src/main.rs | 1 + src/payload.rs | 1 + src/pi_agent.rs | 286 ++++++++++++++++++++++++++++++++++++++++++++++++ src/tags.rs | 21 ++++ 5 files changed, 310 insertions(+) create mode 100644 src/pi_agent.rs diff --git a/src/lib.rs b/src/lib.rs index 709e4e7..0d88cfa 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,6 +2,7 @@ pub mod emit; pub mod log; pub mod opencode; pub mod payload; +pub mod pi_agent; pub mod source; pub mod state; pub mod tags; diff --git a/src/main.rs b/src/main.rs index 4d9b1e4..48b1435 100644 --- a/src/main.rs +++ b/src/main.rs @@ -48,6 +48,7 @@ fn run() -> i32 { let cwd = input.cwd().map(String::from); let env_tags = tags::gather_env_tags(source, cwd.as_deref(), input.agent_version()); + let _lock = state::FileLock::acquire(); let mut global_state = state::load_state(); state::prune_old_sessions(&mut global_state); diff --git a/src/payload.rs b/src/payload.rs index edb587d..dcf2e4e 100644 --- a/src/payload.rs +++ b/src/payload.rs @@ -66,6 +66,7 @@ fn parse_by_source(source: &Source, value: &Value) -> Input { match source { Source::ClaudeCode => parse_claude_code_payload(value), Source::Opencode => parse_opencode_payload(value), + Source::PiAgent => parse_opencode_payload(value), // placeholder; replaced in next commit } } diff --git a/src/pi_agent.rs b/src/pi_agent.rs new file mode 100644 index 0000000..74c0808 --- /dev/null +++ b/src/pi_agent.rs @@ -0,0 +1,286 @@ +use serde_json::{json, Value}; +use std::collections::HashMap; + +pub fn normalize_pi_agent_messages(messages: Vec) -> Vec { + let mut result: Vec = Vec::new(); + let mut pending_tool_results: HashMap = HashMap::new(); + let mut pending_assistant_idx: Option = None; + + for msg in messages { + let role = msg.get("role").and_then(|v| v.as_str()); + + match role { + Some("user") => { + // Flush any pending tool results onto the previous assistant message + if !pending_tool_results.is_empty() { + if let Some(prev_idx) = pending_assistant_idx { + if prev_idx < result.len() { + let mut tr_arr: Vec = Vec::new(); + for (tid, content) in pending_tool_results.drain() { + tr_arr.push(json!({ + "type": "tool_result", + "tool_use_id": tid, + "content": content, + })); + } + result[prev_idx]["message"]["tool_results"] = json!(tr_arr); + } + } + pending_tool_results.clear(); + } + pending_assistant_idx = None; + + let user_content = normalize_user_content(&msg); + if !user_content.is_empty() { + result.push(json!({ + "type": "user", + "message": { + "role": "user", + "content": user_content, + } + })); + } + } + + Some("assistant") => { + // Attach any pending tool results to the previous assistant message + // before emitting the new one + if !pending_tool_results.is_empty() { + if let Some(prev_idx) = pending_assistant_idx { + if prev_idx < result.len() { + let mut tr_arr: Vec = Vec::new(); + for (tid, content) in pending_tool_results.drain() { + tr_arr.push(json!({ + "type": "tool_result", + "tool_use_id": tid, + "content": content, + })); + } + result[prev_idx]["message"]["tool_results"] = json!(tr_arr); + } + } + pending_tool_results.clear(); + } + + let id = msg + .get("id") + .and_then(|v| v.as_str()) + .map(String::from) + .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + + let model = msg + .get("model") + .and_then(|v| v.as_str()) + .unwrap_or("pi"); + + let mut content_parts: Vec = Vec::new(); + let mut has_tool_use = false; + + if let Some(blocks) = msg.get("content").and_then(|v| v.as_array()) { + for block in blocks { + let block_type = block.get("type").and_then(|v| v.as_str()); + match block_type { + Some("text") => { + let text = block.get("text").and_then(|v| v.as_str()).unwrap_or(""); + content_parts.push(json!({ "type": "text", "text": text })); + } + Some("thinking") => { + // Skip thinking blocks + } + Some("toolCall") => { + has_tool_use = true; + let tool_id = block.get("id").and_then(|v| v.as_str()).unwrap_or(""); + let tool_name = block.get("name").and_then(|v| v.as_str()).unwrap_or("unknown"); + let tool_input = block.get("arguments").cloned().unwrap_or(json!({})); + content_parts.push(json!({ + "type": "tool_use", + "id": tool_id, + "name": tool_name, + "input": tool_input, + })); + } + _ => {} + } + } + } + + let assistant_msg_idx = result.len(); + if has_tool_use { + pending_assistant_idx = Some(assistant_msg_idx); + } + + result.push(json!({ + "type": "assistant", + "message": { + "id": id, + "role": "assistant", + "model": model, + "content": content_parts, + } + })); + } + + Some("toolResult") => { + let tool_call_id = msg + .get("toolCallId") + .and_then(|v| v.as_str()) + .unwrap_or(""); + let content = msg.get("content").cloned().unwrap_or(json!([])); + pending_tool_results.insert(tool_call_id.to_string(), content); + } + + // Skip: session, compaction, branch_summary, custom, label, + // model_change, thinking_level_change, session_info, etc. + _ => {} + } + } + + // Flush any remaining pending tool results + if !pending_tool_results.is_empty() { + if let Some(prev_idx) = pending_assistant_idx { + if prev_idx < result.len() { + let mut tr_arr: Vec = Vec::new(); + for (tid, content) in pending_tool_results.drain() { + tr_arr.push(json!({ + "type": "tool_result", + "tool_use_id": tid, + "content": content, + })); + } + result[prev_idx]["message"]["tool_results"] = json!(tr_arr); + } + } + } + + result +} + +fn normalize_user_content(msg: &Value) -> Vec { + match msg.get("content") { + Some(Value::String(s)) => { + vec![json!({ "type": "text", "text": s })] + } + Some(Value::Array(blocks)) => { + blocks + .iter() + .filter_map(|block| { + let block_type = block.get("type").and_then(|v| v.as_str()); + if block_type == Some("text") { + let text = block.get("text").and_then(|v| v.as_str()).unwrap_or(""); + Some(json!({ "type": "text", "text": text })) + } else { + None + } + }) + .collect() + } + _ => vec![], + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn text_only_user_message_string() { + let msgs = vec![json!({ "role": "user", "content": "Hello" })]; + let out = normalize_pi_agent_messages(msgs); + assert_eq!(out.len(), 1); + assert_eq!(out[0]["type"], "user"); + assert_eq!(out[0]["message"]["role"], "user"); + assert_eq!(out[0]["message"]["content"][0]["type"], "text"); + assert_eq!(out[0]["message"]["content"][0]["text"], "Hello"); + } + + #[test] + fn text_only_user_message_array_content() { + let msgs = vec![json!({ + "role": "user", + "content": [{ "type": "text", "text": "Hello" }] + })]; + let out = normalize_pi_agent_messages(msgs); + assert_eq!(out.len(), 1); + assert_eq!(out[0]["type"], "user"); + assert_eq!(out[0]["message"]["role"], "user"); + assert_eq!(out[0]["message"]["content"][0]["type"], "text"); + assert_eq!(out[0]["message"]["content"][0]["text"], "Hello"); + } + + #[test] + fn assistant_message_with_tool_use() { + let msgs = vec![json!({ + "role": "assistant", + "content": [ + { "type": "text", "text": "Let me check" }, + { "type": "toolCall", "id": "tc1", "name": "bash", "arguments": { "command": "ls" } } + ], + "model": "claude" + })]; + let out = normalize_pi_agent_messages(msgs); + assert_eq!(out.len(), 1); + assert_eq!(out[0]["type"], "assistant"); + let content = out[0]["message"]["content"].as_array().unwrap(); + assert_eq!(content[0]["type"], "text"); + assert_eq!(content[0]["text"], "Let me check"); + assert_eq!(content[1]["type"], "tool_use"); + assert_eq!(content[1]["name"], "bash"); + assert_eq!(content[1]["input"]["command"], "ls"); + } + + #[test] + fn tool_result_attached_to_previous_assistant() { + let msgs = vec![ + json!({ + "role": "assistant", + "content": [ + { "type": "text", "text": "Running command" }, + { "type": "toolCall", "id": "tc1", "name": "bash", "arguments": {} } + ], + "model": "claude" + }), + json!({ + "role": "toolResult", + "toolCallId": "tc1", + "toolName": "bash", + "content": [{ "type": "text", "text": "file1.txt" }], + "isError": false + }), + ]; + let out = normalize_pi_agent_messages(msgs); + assert_eq!(out.len(), 2); + let tool_results = out[0]["message"]["tool_results"].as_array().unwrap(); + assert_eq!(tool_results.len(), 1); + assert_eq!(tool_results[0]["tool_use_id"], "tc1"); + } + + #[test] + fn skips_thinking_blocks() { + let msgs = vec![json!({ + "role": "assistant", + "content": [ + { "type": "thinking", "thinking": "internal thoughts" }, + { "type": "text", "text": "Result" } + ], + "model": "claude" + })]; + let out = normalize_pi_agent_messages(msgs); + assert_eq!(out.len(), 1); + let content = out[0]["message"]["content"].as_array().unwrap(); + assert_eq!(content.len(), 1); + assert_eq!(content[0]["type"], "text"); + } + + #[test] + fn skips_unknown_roles() { + let msgs = vec![ + json!({ "role": "session", "data": {} }), + json!({ "role": "compaction", "data": {} }), + json!({ "role": "user", "content": "Hi" }), + ]; + let out = normalize_pi_agent_messages(msgs); + assert_eq!(out.len(), 1); + assert_eq!(out[0]["type"], "user"); + } +} diff --git a/src/tags.rs b/src/tags.rs index 7e613bd..6d2112d 100644 --- a/src/tags.rs +++ b/src/tags.rs @@ -79,6 +79,18 @@ pub fn gather_env_tags(source: Source, cwd: Option<&str>, agent_version: Option< } } + // Pi Agent version probe (only for PiAgent source, if no version provided) + if source == Source::PiAgent && agent_version.is_none() { + if let Ok(output) = Command::new("pi").arg("--version").output() { + if output.status.success() { + let ver = String::from_utf8_lossy(&output.stdout).trim().to_string(); + if !ver.is_empty() { + tags.push(format!("pi-version:{ver}")); + } + } + } + } + tags } @@ -117,4 +129,13 @@ mod tests { let tags = gather_env_tags(Source::Opencode, None, Some("1.2.3")); assert!(tags.contains(&"oc-version:1.2.3".to_string())); } + + #[test] + fn pi_agent_source_uses_pi_agent_tag() { + let tags = gather_env_tags(Source::PiAgent, None, Some("1.0.0")); + assert!(tags.contains(&"pi-agent".to_string())); + assert!(tags.contains(&"pi-version:1.0.0".to_string())); + assert!(!tags.iter().any(|t| t.starts_with("cc-version:"))); + assert!(!tags.iter().any(|t| t.starts_with("oc-version:"))); + } } From 045a9f8699a1278ed5b404193eec9cecc8230399 Mon Sep 17 00:00:00 2001 From: winjer Date: Mon, 11 May 2026 09:49:13 +0100 Subject: [PATCH 04/10] feat: add pi_agent message normalizer Implement normalize_pi_agent_messages converting Pi agent session entries (user, assistant, toolResult roles) into the Claude-format shape expected by turns::build_turns, with thinking-block suppression and tool-result correlation following the opencode.rs pattern. Co-Authored-By: Claude Sonnet 4.6 --- src/pi_agent.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/pi_agent.rs b/src/pi_agent.rs index 74c0808..9268d89 100644 --- a/src/pi_agent.rs +++ b/src/pi_agent.rs @@ -231,6 +231,8 @@ mod tests { #[test] fn tool_result_attached_to_previous_assistant() { + // toolResult entries are not emitted as separate messages; they are + // attached to the previous assistant message's tool_results field. let msgs = vec![ json!({ "role": "assistant", @@ -249,7 +251,8 @@ mod tests { }), ]; let out = normalize_pi_agent_messages(msgs); - assert_eq!(out.len(), 2); + // toolResult is attached to the assistant message, not a separate entry + assert_eq!(out.len(), 1); let tool_results = out[0]["message"]["tool_results"].as_array().unwrap(); assert_eq!(tool_results.len(), 1); assert_eq!(tool_results[0]["tool_use_id"], "tc1"); From bfa87deb5875f531631ba82542b11801ad8aa064 Mon Sep 17 00:00:00 2001 From: winjer Date: Mon, 11 May 2026 09:49:39 +0100 Subject: [PATCH 05/10] feat: add PiAgent Input variant to payload --- src/main.rs | 53 +++++++++++++++++++++++++++++++++++++++++++- src/payload.rs | 60 +++++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 111 insertions(+), 2 deletions(-) diff --git a/src/main.rs b/src/main.rs index 48b1435..2e78ae2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,4 @@ -use code_trace::{emit, log, opencode, payload, state, tags, transcript, turns}; +use code_trace::{emit, log, opencode, payload, pi_agent, state, tags, transcript, turns}; use std::time::Instant; fn run() -> i32 { @@ -173,6 +173,57 @@ fn run() -> i32 { dur.as_secs_f64() )); } + + payload::Input::PiAgent { + session_id: _, + cwd: _, + messages, + agent_version: _, + } => { + if messages.is_empty() { + return 0; + } + + let key = state::state_key(source.as_str(), &session_id, &session_id); + let mut ss = global_state.get(&key).cloned().unwrap_or_default(); + + let normalized = pi_agent::normalize_pi_agent_messages(messages); + let built_turns = turns::build_turns(normalized); + if built_turns.is_empty() { + global_state.insert(key, ss); + state::save_state(&global_state); + return 0; + } + + let mut all_events = Vec::new(); + let mut emitted = 0u32; + for t in &built_turns { + emitted += 1; + let turn_num = ss.turn_count + emitted; + let events = emit::build_ingestion_batch( + &session_id, + turn_num, + t, + std::path::Path::new("pi-agent"), + &env_tags, + source, + ); + all_events.extend(events); + } + + ss.turn_count += emitted; + state::touch(&mut ss); + global_state.insert(key, ss); + state::save_state(&global_state); + + emit::send_batch_fire_and_forget(&config, all_events); + + let dur = start.elapsed(); + log::info(&format!( + "Processed {emitted} turns in {:.2}s (session={session_id})", + dur.as_secs_f64() + )); + } } 0 diff --git a/src/payload.rs b/src/payload.rs index dcf2e4e..14abfea 100644 --- a/src/payload.rs +++ b/src/payload.rs @@ -21,6 +21,12 @@ pub enum Input { messages: Vec, agent_version: Option, }, + PiAgent { + session_id: Option, + cwd: Option, + messages: Vec, + agent_version: Option, + }, } impl Input { @@ -28,6 +34,7 @@ impl Input { match self { Input::ClaudeCode { .. } => Source::ClaudeCode, Input::Opencode { .. } => Source::Opencode, + Input::PiAgent { .. } => Source::PiAgent, } } @@ -35,6 +42,7 @@ impl Input { match self { Input::ClaudeCode { session_id, .. } => session_id.as_deref(), Input::Opencode { session_id, .. } => session_id.as_deref(), + Input::PiAgent { session_id, .. } => session_id.as_deref(), } } @@ -42,6 +50,7 @@ impl Input { match self { Input::ClaudeCode { cwd, .. } => cwd.as_deref(), Input::Opencode { cwd, .. } => cwd.as_deref(), + Input::PiAgent { cwd, .. } => cwd.as_deref(), } } @@ -49,6 +58,7 @@ impl Input { match self { Input::ClaudeCode { .. } => None, Input::Opencode { agent_version, .. } => agent_version.as_deref(), + Input::PiAgent { agent_version, .. } => agent_version.as_deref(), } } } @@ -66,7 +76,7 @@ fn parse_by_source(source: &Source, value: &Value) -> Input { match source { Source::ClaudeCode => parse_claude_code_payload(value), Source::Opencode => parse_opencode_payload(value), - Source::PiAgent => parse_opencode_payload(value), // placeholder; replaced in next commit + Source::PiAgent => parse_pi_agent_payload(value), } } @@ -129,6 +139,35 @@ fn parse_opencode_payload(value: &Value) -> Input { } } +fn parse_pi_agent_payload(value: &Value) -> Input { + let session_id = value + .get("sessionId") + .or_else(|| value.get("session_id")) + .and_then(|v| v.as_str()) + .map(String::from); + + let cwd = value.get("cwd").and_then(|v| v.as_str()).map(String::from); + + let messages = value + .get("messages") + .and_then(|v| v.as_array()) + .cloned() + .unwrap_or_default(); + + let agent_version = value + .get("agentVersion") + .or_else(|| value.get("agent_version")) + .and_then(|v| v.as_str()) + .map(String::from); + + Input::PiAgent { + session_id, + cwd, + messages, + agent_version, + } +} + pub fn read_stdin() -> Value { let mut buf = String::new(); if std::io::Read::read_to_string(&mut std::io::stdin(), &mut buf).is_err() { @@ -195,6 +234,25 @@ mod tests { } } + #[test] + fn parses_pi_agent_payload() { + let v = json!({ + "source": "pi-agent", + "sessionId": "ses_456", + "cwd": "/home/user/project", + "messages": [{"role": "user", "content": "hello"}], + "agentVersion": "1.0.0" + }); + let input = parse_payload(&v); + assert_eq!(input.source(), Source::PiAgent); + assert_eq!(input.session_id(), Some("ses_456")); + assert_eq!(input.agent_version(), Some("1.0.0")); + match &input { + Input::PiAgent { messages, .. } => assert_eq!(messages.len(), 1), + _ => panic!("expected PiAgent input"), + } + } + #[test] fn defaults_to_claude_code_when_source_missing() { let v = json!({"sessionId": "abc"}); From a2f03366e69d209ccc4e9e4831ff9f2003093774 Mon Sep 17 00:00:00 2001 From: winjer Date: Mon, 11 May 2026 10:16:20 +0100 Subject: [PATCH 06/10] test: add end_to_end_pi_agent integration test Adds comprehensive integration test for Pi Agent message flow covering user messages, assistant messages with tool calls, and tool results. Test verifies proper normalization, turn building, and event emission for the pi-agent source. Co-Authored-By: Claude Sonnet 4.5 --- tests/integration_test.rs | 44 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/tests/integration_test.rs b/tests/integration_test.rs index 39b0b78..28056fd 100644 --- a/tests/integration_test.rs +++ b/tests/integration_test.rs @@ -103,3 +103,47 @@ fn end_to_end_opencode_transcript() { assert_eq!(events[0]["body"]["name"], "OpenCode - Turn 1"); assert_eq!(events[0]["body"]["metadata"]["source"], "opencode"); } + +#[test] +fn end_to_end_pi_agent_transcript() { + let msgs = vec![ + json!({ + "role": "user", + "content": "list files" + }), + json!({ + "role": "assistant", + "content": [ + { "type": "text", "text": "Let me check" }, + { "type": "toolCall", "id": "tc1", "name": "bash", "arguments": { "command": "ls" } } + ], + "model": "claude-sonnet-4-5" + }), + json!({ + "role": "toolResult", + "toolCallId": "tc1", + "toolName": "bash", + "content": [{ "type": "text", "text": "file1.txt" }], + "isError": false + }), + ]; + + let normalized = code_trace::pi_agent::normalize_pi_agent_messages(msgs); + let turns = code_trace::turns::build_turns(normalized); + assert_eq!(turns.len(), 1); + + let tags = vec!["pi-agent".to_string()]; + let events = code_trace::emit::build_ingestion_batch( + "sess-pi-1", + 1, + &turns[0], + std::path::Path::new("pi-agent"), + &tags, + code_trace::source::Source::PiAgent, + ); + + assert!(events.len() >= 2); + assert_eq!(events[0]["type"], "trace-create"); + assert!(events[0]["body"]["name"].as_str().unwrap().starts_with("Pi Agent")); + assert_eq!(events[0]["body"]["metadata"]["source"], "pi-agent"); +} From 91b137560dadb244499c1cf05cadd778cf9da88d Mon Sep 17 00:00:00 2001 From: winjer Date: Mon, 11 May 2026 10:16:34 +0100 Subject: [PATCH 07/10] feat: add pi extension installation to install.sh Add support for installing Pi Agent extension alongside the existing OpenCode plugin support. Includes --pi/-p flag, pi detection, and automatic or prompted installation to ~/.pi/agent/extensions. Co-Authored-By: Claude Sonnet 4.5 --- install.sh | 48 ++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 46 insertions(+), 2 deletions(-) diff --git a/install.sh b/install.sh index d03d709..d3a126c 100755 --- a/install.sh +++ b/install.sh @@ -6,6 +6,7 @@ BINARY="code-trace" INSTALL_DIR="${HOME}/.local/bin" SETTINGS_FILE="${HOME}/.claude/settings.json" OPENCODE_PLUGIN_DIR="${HOME}/.config/opencode/plugins" +PI_EXTENSION_DIR="${HOME}/.pi/agent/extensions" # Parse flags INSTALL_OPENCODE=false @@ -13,10 +14,19 @@ if [ "${1:-}" = "--opencode" ] || [ "${1:-}" = "-o" ]; then INSTALL_OPENCODE=true fi +INSTALL_PI=false +if [ "${1:-}" = "--pi" ] || [ "${1:-}" = "-p" ]; then + INSTALL_PI=true +fi + detect_opencode() { [ -d "${HOME}/.config/opencode" ] || [ -f "${HOME}/.config/opencode/opencode.json" ] } +detect_pi() { + [ -d "${HOME}/.pi/agent" ] +} + # Detect platform OS="$(uname -s | tr '[:upper:]' '[:lower:]')" ARCH="$(uname -m)" @@ -85,6 +95,13 @@ elif [ -f "${SCRIPT_DIR}/../plugin/code-trace.ts" ]; then PLUGIN_SRC="$(cd "${SCRIPT_DIR}/../plugin" && pwd)/code-trace.ts" fi +PI_PLUGIN_SRC="" +if [ -f "${SCRIPT_DIR}/plugin/pi-agent/code-trace.ts" ]; then + PI_PLUGIN_SRC="${SCRIPT_DIR}/plugin/pi-agent/code-trace.ts" +elif [ -f "${SCRIPT_DIR}/../plugin/pi-agent/code-trace.ts" ]; then + PI_PLUGIN_SRC="$(cd "${SCRIPT_DIR}/../plugin/pi-agent" && pwd)/code-trace.ts" +fi + # Register Claude Code hook HOOK_ENTRY='{"type":"command","command":"code-trace"}' @@ -149,6 +166,20 @@ install_opencode_plugin() { echo "Installed OpenCode plugin to ${OPENCODE_PLUGIN_DIR}/code-trace.ts" } +# Install Pi Agent extension +install_pi_extension() { + if [ -z "${PI_PLUGIN_SRC}" ] || [ ! -f "${PI_PLUGIN_SRC}" ]; then + echo "" + echo "Note: Pi extension source not found at ${PI_PLUGIN_SRC}. Skipping Pi Agent extension install." + echo "You can manually copy plugin/pi-agent/code-trace.ts to ${PI_EXTENSION_DIR}/" + return + fi + + mkdir -p "${PI_EXTENSION_DIR}" + cp "${PI_PLUGIN_SRC}" "${PI_EXTENSION_DIR}/code-trace.ts" + echo "Installed Pi Agent extension to ${PI_EXTENSION_DIR}/code-trace.ts" +} + if [ "${INSTALL_OPENCODE}" = true ]; then install_opencode_plugin elif detect_opencode; then @@ -162,9 +193,22 @@ elif detect_opencode; then fi fi +if [ "${INSTALL_PI}" = true ]; then + install_pi_extension +elif detect_pi; then + echo "" + echo "Pi Agent detected. Install the code-trace extension?" + echo " ${PI_EXTENSION_DIR}/code-trace.ts" + echo "" + read -p "Install Pi Agent extension? [y/N] " -r + if [[ "${REPLY}" =~ ^[Yy]$ ]]; then + install_pi_extension + fi +fi + echo "" echo "Done! To enable tracing, add to your project's .claude/settings.local.json (Claude Code)" -echo "or set environment variables for OpenCode:" +echo "or set environment variables for OpenCode and Pi Agent:" echo "" cat << 'EOF' { @@ -176,4 +220,4 @@ cat << 'EOF' } EOF echo "" -echo "For OpenCode, set these environment variables in your shell profile." +echo "For OpenCode and Pi Agent extensions, set these environment variables in your shell profile." From 893ef7c2ff8cc0c9a5779b9fe3f278541300092f Mon Sep 17 00:00:00 2001 From: winjer Date: Mon, 11 May 2026 10:16:36 +0100 Subject: [PATCH 08/10] docs: update CLAUDE.md for pi-agent source Add Pi Agent as a third input source in the architecture section. Document the pi_agent normalizer and pi extension in key files. Co-Authored-By: Claude Sonnet 4.5 --- CLAUDE.md | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 359c554..f26e881 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -2,22 +2,25 @@ This project sends Claude Code and OpenCode session traces to Langfuse for obser ## Architecture -- Single Rust binary that supports two input sources +- Single Rust binary that supports three input sources - Claude Code: invoked as a Stop hook, reads transcript JSONL - OpenCode: invoked by a TypeScript plugin, receives messages via stdin +- Pi: invoked by a TypeScript extension, receives session entries via stdin ## Key files - `src/main.rs` — entry point, dispatches on `Input` enum -- `src/source.rs` — `Source` enum (ClaudeCode | Opencode) +- `src/source.rs` — `Source` enum (ClaudeCode | Opencode | PiAgent) - `src/payload.rs` — parses stdin into `Input` enum - `src/opencode.rs` — normalizes OpenCode SDK message format to Claude-format Values +- `src/pi_agent.rs` — normalizes Pi session entry format to Claude-format Values - `src/transcript.rs` — reads Claude Code JSONL transcript - `src/turns.rs` — groups messages into user/assistant/tool turns - `src/emit.rs` — builds Langfuse ingestion batch - `src/tags.rs` — gathers env tags (repo, branch, user, host, os, agent version) - `src/state.rs` — persisted cursor state per session - `src/log.rs` — logging to `~/.local/share/code-trace/` +- `plugin/pi-agent/code-trace.ts` — pi extension ## State location From dfab424b451aabf096a5d8b1d086989fe0a304ac Mon Sep 17 00:00:00 2001 From: winjer Date: Mon, 11 May 2026 10:18:14 +0100 Subject: [PATCH 09/10] feat: add pi-agent TypeScript extension Co-Authored-By: Claude Sonnet 4.6 --- plugin/pi-agent/code-trace.ts | 90 +++++++++++++++++++++++++++++++++++ 1 file changed, 90 insertions(+) create mode 100644 plugin/pi-agent/code-trace.ts diff --git a/plugin/pi-agent/code-trace.ts b/plugin/pi-agent/code-trace.ts new file mode 100644 index 0000000..a637b5a --- /dev/null +++ b/plugin/pi-agent/code-trace.ts @@ -0,0 +1,90 @@ +import { spawn } from "node:child_process"; +import { readFileSync, writeFileSync, existsSync, mkdirSync } from "node:fs"; +import { join, dirname } from "node:path"; + +interface CursorEntry { + lastIndex: number; +} + +interface CursorStore { + [sessionId: string]: CursorEntry; +} + +function getCursorPath(): string { + const home = process.env.HOME ?? process.env.USERPROFILE ?? "~"; + return join(home, ".local", "share", "code-trace", "pi_agent_cursor.json"); +} + +function ensureDir(path: string): void { + const dir = dirname(path); + if (!existsSync(dir)) { + mkdirSync(dir, { recursive: true }); + } +} + +function loadCursor(): CursorStore { + const path = getCursorPath(); + try { + if (existsSync(path)) { + return JSON.parse(readFileSync(path, "utf-8")); + } + } catch { + // ignore + } + return {}; +} + +function saveCursor(store: CursorStore): void { + const path = getCursorPath(); + ensureDir(path); + writeFileSync(path, JSON.stringify(store, null, 2)); +} + +export default function (pi: any) { + pi.on("agent_end", async (event: any, ctx: any) => { + if (process.env.TRACE_TO_LANGFUSE?.toLowerCase() !== "true") return; + + const allEntries: any[] = ctx.sessionManager.getEntries(); + const sessionId: string | undefined = allEntries[0]?.id; + if (!sessionId) return; + + const cursor = loadCursor(); + const startIndex = cursor[sessionId]?.lastIndex ?? 0; + + if (allEntries.length <= startIndex) return; + + const newEntries = allEntries.slice(startIndex); + if (newEntries.length === 0) return; + + let agentVersion: string | undefined; + try { + const result = await pi.exec("pi", ["--version"], { signal: ctx.signal }); + agentVersion = result.stdout?.trim() || undefined; + } catch { + // ignore + } + + const payload = { + source: "pi-agent", + sessionId, + cwd: ctx.cwd, + messages: newEntries, + agentVersion, + }; + + const binPath = process.env.CODE_TRACE_BIN ?? "code-trace"; + + const child = spawn(binPath, [], { + stdio: ["pipe", "ignore", "ignore"], + detached: true, + shell: true, + }); + + child.stdin?.end(JSON.stringify(payload), "utf-8"); + child.stdin?.destroy(); + child.unref(); + + cursor[sessionId] = { lastIndex: allEntries.length }; + saveCursor(cursor); + }); +} From 839c3e51f81475e3282351dc299a89807fc8cd43 Mon Sep 17 00:00:00 2001 From: winjer Date: Mon, 11 May 2026 10:23:27 +0100 Subject: [PATCH 10/10] refactor: remove redundant clear() after drain() in pi_agent normalizer Co-Authored-By: Claude Sonnet 4.6 --- src/pi_agent.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/pi_agent.rs b/src/pi_agent.rs index 9268d89..5ab7bc9 100644 --- a/src/pi_agent.rs +++ b/src/pi_agent.rs @@ -26,7 +26,6 @@ pub fn normalize_pi_agent_messages(messages: Vec) -> Vec { result[prev_idx]["message"]["tool_results"] = json!(tr_arr); } } - pending_tool_results.clear(); } pending_assistant_idx = None; @@ -59,7 +58,6 @@ pub fn normalize_pi_agent_messages(messages: Vec) -> Vec { result[prev_idx]["message"]["tool_results"] = json!(tr_arr); } } - pending_tool_results.clear(); } let id = msg