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..f26e881 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,11 +1,49 @@ -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 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 -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 | 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 + +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..d3a126c 100755 --- a/install.sh +++ b/install.sh @@ -5,6 +5,27 @@ 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" +PI_EXTENSION_DIR="${HOME}/.pi/agent/extensions" + +# Parse flags +INSTALL_OPENCODE=false +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:]')" @@ -65,15 +86,29 @@ 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 + +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"}' 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 +119,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 +152,63 @@ 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" +} + +# 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 + 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 + +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:" +echo "Done! To enable tracing, add to your project's .claude/settings.local.json (Claude Code)" +echo "or set environment variables for OpenCode and Pi Agent:" echo "" cat << 'EOF' { @@ -130,3 +219,5 @@ cat << 'EOF' } } EOF +echo "" +echo "For OpenCode and Pi Agent extensions, 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/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); + }); +} 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..0d88cfa 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,9 @@ 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; 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..2e78ae2 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, pi_agent, 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,199 @@ 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; - } + let cwd = input.cwd().map(String::from); + let env_tags = tags::gather_env_tags(source, cwd.as_deref(), input.agent_version()); - // Gather tags - let env_tags = tags::gather_env_tags(hook.cwd.as_deref()); - // 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; - } - - 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); + 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() + )); + } + + 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() + )); + } } - 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..14abfea 100644 --- a/src/payload.rs +++ b/src/payload.rs @@ -1,14 +1,86 @@ 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, + }, + PiAgent { + 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, + Input::PiAgent { .. } => Source::PiAgent, + } + } + + 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(), + Input::PiAgent { 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(), + Input::PiAgent { cwd, .. } => cwd.as_deref(), + } + } + + pub fn agent_version(&self) -> Option<&str> { + match self { + Input::ClaudeCode { .. } => None, + Input::Opencode { agent_version, .. } => agent_version.as_deref(), + Input::PiAgent { 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), + Source::PiAgent => parse_pi_agent_payload(value), + } +} + +fn parse_claude_code_payload(value: &Value) -> Input { let session_id = value .get("sessionId") .or_else(|| value.get("session_id")) @@ -29,19 +101,73 @@ 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, + } +} + +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, } } -/// 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 +188,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 +199,71 @@ 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 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"}); + 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/pi_agent.rs b/src/pi_agent.rs new file mode 100644 index 0000000..5ab7bc9 --- /dev/null +++ b/src/pi_agent.rs @@ -0,0 +1,287 @@ +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_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); + } + } + } + + 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() { + // 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", + "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); + // 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"); + } + + #[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/source.rs b/src/source.rs new file mode 100644 index 0000000..99cf316 --- /dev/null +++ b/src/source.rs @@ -0,0 +1,52 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] +pub enum Source { + #[default] + ClaudeCode, + Opencode, + PiAgent, +} + +impl Source { + pub fn agent_tag(&self) -> &'static str { + match self { + Source::ClaudeCode => "claude-code", + Source::Opencode => "opencode", + Source::PiAgent => "pi-agent", + } + } + + pub fn version_tag_prefix(&self) -> &'static str { + match self { + Source::ClaudeCode => "cc-version", + Source::Opencode => "oc-version", + Source::PiAgent => "pi-version", + } + } + + pub fn trace_name_prefix(&self) -> &'static str { + match self { + Source::ClaudeCode => "Claude Code", + Source::Opencode => "OpenCode", + Source::PiAgent => "Pi Agent", + } + } + + pub fn as_str(&self) -> &'static str { + match self { + Source::ClaudeCode => "claude-code", + Source::Opencode => "opencode", + Source::PiAgent => "pi-agent", + } + } + + 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), + "pi-agent" | "pi_agent" | "pi" => Some(Source::PiAgent), + _ => 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..6d2112d 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,38 @@ 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}")); + } + } + } + } + + // 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}")); + } } } } @@ -65,20 +99,43 @@ 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())); + } + + #[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:"))); + } } diff --git a/tests/integration_test.rs b/tests/integration_test.rs index 7bd835b..28056fd 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,89 @@ 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"); +} + +#[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"); +}