Skip to content

Index now supports Codex and Pi traces#49

Merged
dacorvo merged 8 commits into
mainfrom
agent-trace-formats
Jul 2, 2026
Merged

Index now supports Codex and Pi traces#49
dacorvo merged 8 commits into
mainfrom
agent-trace-formats

Conversation

@dacorvo

@dacorvo dacorvo commented Jul 2, 2026

Copy link
Copy Markdown
Collaborator

In addition to Claude code traces and Hugging Face Hub parquet, this adds support for the Codex and Pi traces format.

dacorvo and others added 7 commits July 2, 2026 19:54
Factor the format-agnostic JSONL machinery (recursive file walk,
whole-file read+parse, tool-name back-fill) out of claude_traces into a
new jsonl.rs so the coming codex/pi parsers reuse it. Add harness.rs
owning the Harness enum + detection (a known session dir wins, else a
first-line `type` sniff) and the known-roots helper — scaffolding for
multi-harness ingestion, not yet wired to a parser.

No behavior change: claude_traces routes through the shared helpers and
its tests pass unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add a small, secret-scanned real Claude Code session under tests/fixtures/
and a deterministic (no-network) test that parses it, guarding claude_traces
against real-format drift the synthetic unit tests can't see: real skip-line
types (queue-operation), tool-name correlation over a real tool_use/tool_result
chain, and turn_uuid stability across a re-parse.

The fixture is a trimmed slice of a public Hub agent-traces dataset with
over-long fields shortened; scanned clean with trufflehog before committing.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Parse Codex native rollout transcripts (~/.codex/sessions/rollout-*.jsonl)
directly into the shared Turn/Block model. Keep only `response_item` lines;
map payload.type: message->text, reasoning->thinking (from summary; the real
reasoning in encrypted_content is opaque and ignored), function_call /
custom_tool_call->tool_use (tool_use_id=call_id, args verbatim from `arguments`
or the raw `input`), *_output->tool_result; skip session_meta / event_msg /
turn_context / tool_schema. Each retained item is its own turn with a stable
turn_uuid=session_id-seq so incremental re-index dedups correctly; tool names
correlate via the shared jsonl::backfill_tool_names.

Add a real Codex session fixture (trimmed, trufflehog-clean) and extend the
real-session parse test to cover it. Not yet wired to `funes index` (Step 4)
and does not stamp the harness facet (Steps 5-6).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Parse pi native session transcripts (~/.pi/agent/sessions/*.jsonl) directly
into the shared Turn/Block model. Conversation lives in `type:"message"` lines
(a nested `message` with a typed content-part list); control events
(session/model_change/thinking_level_change) carry no turn. Map the parts:
thinking->thinking, text->text (blank whitespace dropped), toolCall->tool_use
(arguments serialized compactly); a `toolResult` message->tool_result. pi names
its tool result inline (toolName), so there is no back-fill pass. turn_uuid is
the native line `id` (stable across re-reads); a result's role is normalized to
`tool` to match codex_traces.

Add a real pi session fixture (trimmed, trufflehog-clean) and extend the
real-session parse test to cover it. Not yet wired to `funes index` (Step 4)
and does not stamp the harness facet (Steps 5-7).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Wire the three native parsers behind a detected harness and generalize
`funes index` to a multi-root command.

- source: `open_with_harness` detects a JSONL tree's harness (a known session
  dir wins, else sniff the first record via the new `jsonl::first_record`) or
  takes a `--harness` override; JsonlTree carries the harness and `read()`
  dispatches to claude/codex/pi. The dispatch layer stays harness-agnostic —
  Codex derives its own session id from the `session_meta` line inside
  `codex_traces` (one pass, no extra per-file read), rather than a helper here.
- index: `run_index_roots` indexes several roots sharing one store, embedder,
  and state.json (keyed by absolute path, so cross-root incremental works);
  `run_index` stays as a single-root convenience.
- main: `funes index [PATH|REPO] [--harness]` — an existing path wins, else an
  `<org>/<repo>`/`hf://` ref (bails until the remote source lands), else index
  every present known harness dir. `hub::is_remote_shorthand` made pub.

Repo-ref indexing and the harness facet are still to come.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Record which coding agent produced each session and make it a first-class
recall facet.

- Add harness:String to Turn and Chunk, and to the lance schema as the LAST
  column (after `vector`) so a later `add_columns` migration lines up.
- Set it at every construction site: each JSONL parser stamps its own constant
  (claude_traces->claude_code, codex_traces->codex, pi_traces->pi), while
  hf_traces reads it per-row from the parquet `harness` column (a dataset can
  mix harnesses).
- Recall: a `--harness` filter (applied as a prefilter via build_where) and the
  harness shown in each hit; the MCP recall tool gains the same filter. Reads
  tolerate a pre-migration store — an absent column reads as "".

Existing stores read fine now; appending to one still needs the column added,
which the coming `funes migrate` handles.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Resolve a Hub trace dataset's `refs/convert/parquet` branch, download each
shard whole-file through the existing hf-hub cache, and index it through the
same units→skip→read→embed→write pipeline as the local sources.

Resolution uses hf-hub's public `list_refs` + `list_tree` (find the `parquet`
convert ref, list its `.parquet` shards) — not the datasets-server `/parquet`
HTTP endpoint, which would add a `reqwest` dep funes doesn't carry. Each shard
is a unit signed with the convert-branch commit oid, so an unchanged repo is
skipped without re-downloading; a changed repo re-reads and chunk-id dedup
keeps already-stored rows a no-op.

Resolve + download happen in the async `open_remote` constructor so
`TraceSource::read` stays sync (no `block_on` on a Tokio worker). The shared
index body is factored into `index_sources`, so local and remote share it.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR extends funes index / funes recall to support additional trace “harnesses” (Codex and Pi) alongside existing Claude Code traces, and adds first-class support for indexing Hugging Face “agent-traces” datasets via their refs/convert/parquet shards. It also introduces a harness facet that is stored in the index and can be used to filter recall results.

Changes:

  • Add Codex and Pi JSONL parsers plus shared JSONL utilities (walk/read/first-record/tool-name backfill), and record a per-turn harness.
  • Add remote Hub dataset indexing: resolve refs/convert/parquet, download shards into cache, and index them incrementally.
  • Extend recall to project/use the harness column, including a regression test ensuring recall works against stores built before the column existed.

Reviewed changes

Copilot reviewed 30 out of 30 changed files in this pull request and generated 4 comments.

Show a summary per file
File Description
tests/remote_repo_source.rs New gated live test for resolving/downloading/reading a real Hub refs/convert/parquet dataset shard.
tests/remote_recall.rs Update recall invocation to match the new recall(..., harness) signature.
tests/remote_cache.rs Update recall invocation to match the new recall(..., harness) signature.
tests/recall_no_harness_column.rs New regression test: recall must work when the store lacks the harness column.
tests/push_round_trip.rs Update recall invocation to match the new recall(..., harness) signature.
tests/parse_real_sessions.rs New deterministic fixture-based end-to-end parsing tests for Codex/Pi/Claude formats.
tests/index_recall.rs Update recall invocation to match the new recall(..., harness) signature.
tests/hello_world.rs Update recall invocation to match the new recall(..., harness) signature.
tests/fixtures/pi_session.jsonl New real-format Pi fixture JSONL.
tests/fixtures/codex_session.jsonl New real-format Codex fixture JSONL.
tests/fixtures/claude_session.jsonl New real-format Claude fixture JSONL.
src/trace.rs Add Turn.harness field to tag turns by producing agent.
src/source.rs Add harness-aware JSONL source handling and a new remote-parquet trace source (open_remote).
src/recall.rs Add harness projection/filtering and compatibility logic for stores without the column.
src/push.rs Update test data to populate Turn.harness.
src/pi_traces.rs New Pi JSONL parser into the shared turn/block model.
src/mcp.rs Extend MCP recall request schema with optional harness filter.
src/main.rs Extend CLI: index from multiple known harness roots, index remote Hub datasets, and add --harness options.
src/lib.rs Export new modules (codex_traces, pi_traces, jsonl, harness).
src/jsonl.rs New shared JSONL utilities (walk/read/first-record/tool backfill).
src/index.rs Index multiple roots and remote datasets; add harness column to dataset schema and batch writes.
src/hub.rs Make is_remote_shorthand public for CLI routing decisions.
src/hf_traces.rs Read per-row harness from parquet (defaulting to empty when absent).
src/hf_dataset.rs Add Hub parquet resolution (refs/convert/parquet) and shard downloading helpers.
src/hello.rs Populate harness in built-in hello dataset rows.
src/harness.rs New harness enum + parsing + known-root detection.
src/codex_traces.rs New Codex rollout JSONL parser into the shared turn/block model.
src/claude_traces.rs Refactor Claude parser to use shared JSONL utilities and set harness.
src/chunk.rs Thread harness through chunks and chunk extraction from dataset batches/turns.
benchmark/bench_recall.rs Update recall invocation to match the new recall(..., harness) signature.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread src/source.rs
Comment thread src/source.rs Outdated
Comment thread src/source.rs
Comment thread src/recall.rs
Four review comments, all valid on inspection:

- Remote `limit` was overloaded: `open_remote` downloaded every shard then
  truncated in `units()`, and `read()` reused the same value as a per-shard
  *session* cap — silently dropping sessions within a shard. Rename it to
  `max_shards`, apply it once (truncate the shard list before downloading), and
  read all sessions in each kept shard. The gated live read now returns every
  session in the shard, not one.

- `detect_harness` walked and sorted the whole session tree just to sniff the
  first record, even when the directory tail already identified the harness —
  contradicting the documented "known dir wins". Short-circuit on
  `Harness::from_known_dir` before touching the filesystem.

- `recall --harness` against a store built before the harness column existed
  failed deep inside Lance with an opaque schema error. Detect the missing
  column and refuse with a clear message; covered by an added assertion in the
  existing regression test.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@dacorvo
dacorvo merged commit 2d5d782 into main Jul 2, 2026
7 checks passed
@dacorvo
dacorvo deleted the agent-trace-formats branch July 2, 2026 19:45
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants