Add agent-trace format support

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: yanpgwang

CHANGELOG.md

@@ -12,6 +12,7 @@
   - `normalize`
   - `diagnose`
 - Cursor and Claude Code transcript JSONL ingestion.
+- `pagarsky/agent-trace` sample ingestion via `--format agent-trace`.
 - Agent-readable diagnosis reports with input scope, confidence, limitations, issue codes, and recommendations.
 - Context diff evidence:
   - turn-to-turn added/removed/retained token summary

README.md

@@ -97,6 +97,12 @@ npx langfuse-cli api traces get <trace-id> --fields core,io,observations --json
   | context-profiler diagnose - --format langfuse --json
 ```
 
+Analyze a public academic agent trajectory:
+
+```bash
+context-profiler diagnose examples/agent-trace/sample.json --format agent-trace --json
+context-profiler analyze examples/agent-trace/sample.json --format agent-trace --html report.html
+```
 
 Generate an interactive report:
 
@@ -199,6 +205,16 @@ The HTML report is self-contained and keeps the existing profiler style:
 }
 ```
 
+Academic trajectory sample:
+
+```text
+context-profiler analyze examples/agent-trace/sample.json --format agent-trace --html report.html
+
+Total Input: 11.7K
+Messages (assistant): 10.7K
+Tool: python_interpreter 2.9K
+Warnings: Content duplication 2.3K redundant tokens
+```
 
 ## Examples
 

examples/README.md

@@ -75,6 +75,28 @@ context-profiler analyze trace.json --format langfuse --html /tmp/context-profil
 
 `context-profiler` does not fetch Langfuse data. Use Langfuse's own CLI/API or any other tool to obtain the trace first.
 
+## Academic AgentTrace
+
+This repository includes a small sample from [`pagarsky/agent-trace`](https://huggingface.co/datasets/pagarsky/agent-trace):
+
+```bash
+PYTHONPATH=src uv run context-profiler diagnose examples/agent-trace/sample.json --format agent-trace --json
+PYTHONPATH=src uv run context-profiler analyze examples/agent-trace/sample.json --format agent-trace --html /tmp/context-profiler-agent-trace-demo.html
+```
+
+The sample is a multi-step MBPP trajectory with:
+
+- 11 LLM steps
+- 10 `python_interpreter` tool spans
+- repeated code passed into tool calls
+- turn-to-turn growth suitable for recording the timeline
+
+Expected findings:
+
+- `REPEATED_CONTENT_BLOCK`
+- `REPEATED_TOOL_INPUT` on `python_interpreter.code`
+- `large_addition` diff hints
+
 ## Adapting Other Formats
 
 Agents should prefer the built-in format registry and schema:

examples/agent-trace/sample.json

@@ -0,0 +1,169 @@
+"""Adapter for pagarsky/agent-trace style academic trajectories."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from context_profiler.models import APIRequest, BlockType, ContentBlock, Message, Role, Session
+from context_profiler.token_utils import count_tokens
+
+
+def is_agent_trace(data: dict[str, Any]) -> bool:
+    return (
+        ("llm_steps" in data or "llm_steps_json" in data)
+        and ("spans" in data or "spans_json" in data)
+    )
+
+
+def load_agent_trace_session(path: Path) -> Session:
+    with open(path, encoding="utf-8") as f:
+        data = json.load(f)
+    return parse_agent_trace(data, source=str(path))
+
+
+def parse_agent_trace(data: dict[str, Any], source: str = "") -> Session:
+    llm_steps = _json_field(data, "llm_steps")
+    spans = _json_field(data, "spans")
+    metadata = _json_field(data, "metadata") if ("metadata" in data or "metadata_json" in data) else {}
+
+    messages: list[Message] = []
+    requests: list[APIRequest] = []
+
+    prompt = data.get("prompt")
+    if prompt:
+        messages.append(Message(
+            role=Role.USER,
+            blocks=[_text_block(prompt)],
+            index=len(messages),
+        ))
+
+    span_idx = 0
+    for step in llm_steps:
+        assistant_blocks = _assistant_blocks(step)
+        if assistant_blocks:
+            messages.append(Message(role=Role.ASSISTANT, blocks=assistant_blocks, index=len(messages)))
+
+        tool_calls = step.get("tool_calls") or []
+        for call in tool_calls:
+            span = spans[span_idx] if span_idx < len(spans) else None
+            if span is not None:
+                span_idx += 1
+                messages.append(_tool_result_message(span, len(messages)))
+
+        requests.append(_snapshot(messages, len(requests), data))
+
+    return Session(
+        requests=requests,
+        metadata={
+            "source": source,
+            "source_format": "agent-trace",
+            "trace_id": data.get("trace_id"),
+            "dataset_name": data.get("dataset_name") or metadata.get("dataset_name"),
+            "task_id": data.get("task_id") or metadata.get("task_id"),
+            "model": data.get("model") or metadata.get("model_family"),
+            "llm_step_count": len(llm_steps),
+            "tool_span_count": len(spans),
+        },
+    )
+
+
+def _json_field(data: dict[str, Any], name: str) -> Any:
+    if name in data:
+        return data[name]
+    raw = data.get(f"{name}_json")
+    if raw is None:
+        return [] if name in {"llm_steps", "spans"} else {}
+    return json.loads(raw)
+
+
+def _assistant_blocks(step: dict[str, Any]) -> list[ContentBlock]:
+    blocks = []
+    reasoning = step.get("reasoning_content")
+    model_output = step.get("model_output")
+    text_parts = [part for part in (reasoning, model_output) if part]
+    if text_parts:
+        blocks.append(_text_block("\n\n".join(text_parts)))
+
+    for idx, call in enumerate(step.get("tool_calls") or []):
+        name = call.get("name") or "unknown"
+        arguments = call.get("arguments") or {}
+        text = json.dumps(arguments, ensure_ascii=False)
+        blocks.append(ContentBlock(
+            block_type=BlockType.TOOL_USE,
+            text=text,
+            token_count=count_tokens(text),
+            tool_name=name,
+            tool_call_id=f"{step.get('step_id', 'step')}:tool:{idx}",
+            tool_input=arguments if isinstance(arguments, dict) else {"value": arguments},
+        ))
+    return blocks
+
+
+def _tool_result_message(span: dict[str, Any], index: int) -> Message:
+    output = span.get("tool_output")
+    if not isinstance(output, str):
+        output = json.dumps(output, ensure_ascii=False)
+    block = ContentBlock(
+        block_type=BlockType.TOOL_RESULT,
+        text=output or "",
+        token_count=count_tokens(output or ""),
+        tool_name=span.get("tool_name"),
+        tool_call_id=span.get("span_id"),
+    )
+    return Message(role=Role.TOOL, blocks=[block], index=index)
+
+
+def _text_block(text: str) -> ContentBlock:
+    return ContentBlock(
+        block_type=BlockType.TEXT,
+        text=text,
+        token_count=count_tokens(text),
+    )
+
+
+def _snapshot(messages: list[Message], request_index: int, data: dict[str, Any]) -> APIRequest:
+    return APIRequest(
+        messages=list(messages),
+        tools=[],
+        model=data.get("model", "unknown"),
+        request_index=request_index,
+        source_format="agent-trace",
+        raw_input={
+            "model": data.get("model", "unknown"),
+            "messages": [_message_to_raw(message) for message in messages],
+            "tools": [],
+            "source_format": "agent-trace",
+        },
+    )
+
+
+def _message_to_raw(message: Message) -> dict[str, Any]:
+    content = []
+    tool_calls = []
+    for block in message.blocks:
+        if block.block_type == BlockType.TEXT:
+            content.append({"type": "text", "text": block.text})
+        elif block.block_type == BlockType.TOOL_USE:
+            tool_calls.append({
+                "id": block.tool_call_id,
+                "type": "function",
+                "function": {
+                    "name": block.tool_name or "unknown",
+                    "arguments": block.text,
+                },
+            })
+        elif block.block_type == BlockType.TOOL_RESULT:
+            content.append({
+                "type": "tool_result",
+                "tool_use_id": block.tool_call_id,
+                "content": block.text,
+            })
+
+    raw: dict[str, Any] = {"role": message.role.value, "content": content}
+    if len(content) == 1 and content[0].get("type") == "text":
+        raw["content"] = content[0]["text"]
+    if tool_calls:
+        raw["tool_calls"] = tool_calls
+    return raw

src/context_profiler/adapters/agent_trace_adapter.py

@@ -8,6 +8,7 @@
 from typing import Any
 
 from context_profiler.adapters.auto_detect import detect_adapter
+from context_profiler.adapters.agent_trace_adapter import is_agent_trace, load_agent_trace_session
 from context_profiler.adapters.langfuse_adapter import is_langfuse_trace, parse_langfuse_trace
 from context_profiler.adapters.transcript_adapter import (
     TRANSCRIPT_FORMATS,
@@ -95,6 +96,20 @@ def try_load_langfuse(path: Path) -> Session | None:
     return None
 
 
+def try_load_agent_trace(path: Path) -> Session | None:
+    """Try to load as an AgentTrace trajectory file."""
+    if not path.is_file() or path.suffix not in (".json",):
+        return None
+    try:
+        with open(path) as f:
+            data = json.load(f)
+        if is_agent_trace(data):
+            return load_agent_trace_session(path)
+    except (json.JSONDecodeError, KeyError):
+        pass
+    return None
+
+
 def load_multi_trace_session(paths: list[Path], format_hint: str | None = None) -> Session:
     """Load multiple Langfuse trace files and merge into a single Session.
 
@@ -144,13 +159,20 @@ def load_session(path: Path, format_hint: str | None = None) -> Session:
     if format_hint == "langfuse":
         return load_langfuse_trace(path)
 
+    if format_hint == "agent-trace":
+        return load_agent_trace_session(path)
+
     if format_hint in TRANSCRIPT_FORMATS:
         return load_transcript_session(path, source_format=format_hint)
 
     langfuse_session = try_load_langfuse(path)
     if langfuse_session is not None:
         return langfuse_session
 
+    agent_trace_session = try_load_agent_trace(path)
+    if agent_trace_session is not None:
+        return agent_trace_session
+
     requests: list[APIRequest] = []
 
     if path.is_dir():

src/context_profiler/profiler.py

@@ -7,6 +7,7 @@
 from typing import Any
 
 from context_profiler.adapters.auto_detect import detect_adapter
+from context_profiler.adapters.agent_trace_adapter import is_agent_trace
 from context_profiler.adapters.langfuse_adapter import is_langfuse_trace, parse_langfuse_trace
 from context_profiler.adapters.transcript_adapter import TRANSCRIPT_FORMATS
 from context_profiler.models import APIRequest, Session
@@ -58,6 +59,8 @@ def validate_input(path: str, format_hint: str | None = None) -> dict[str, Any]:
             if not session.requests:
                 return _invalid_langfuse_result(session)
             return _valid_result("langfuse")
+        if isinstance(data, dict) and is_agent_trace(data):
+            return _valid_result("agent-trace")
         if isinstance(data, dict):
             adapter = detect_adapter(data) if format_hint is None else None
             return _valid_result(format_hint or _adapter_format_name(adapter))
@@ -142,6 +145,8 @@ def normalize_input(path: str, format_hint: str | None = None) -> dict[str, Any]
 
     if isinstance(data, dict) and is_langfuse_trace(data):
         session = parse_langfuse_trace(data)
+    elif isinstance(data, dict) and is_agent_trace(data):
+        session = load_session(Path(path), format_hint="agent-trace")
     elif path != "-":
         session = load_session(Path(path), format_hint=format_hint)
     elif isinstance(data, dict):

src/context_profiler/validation.py

@@ -286,3 +286,25 @@ def test_skill_distribution_manifests():
     marketplace_data = json.loads(claude_marketplace.read_text())
     assert marketplace_data["plugins"][0]["name"] == "context-profiler"
     assert marketplace_data["plugins"][0]["skills"] == ["./skills/analyze-agent-context"]
+
+
+def test_diagnose_agent_trace_sample_json():
+    sample = Path(__file__).parents[1] / "examples" / "agent-trace" / "sample.json"
+    runner = CliRunner()
+    result = runner.invoke(main, ["diagnose", str(sample), "--format", "agent-trace", "--json"])
+    assert result.exit_code == 0
+    data = json.loads(result.output)
+    assert data["analysis_scope"]["format"] == "agent-trace"
+    assert data["mode"] == "session"
+    assert data["diff_summary"]["transition_count"] > 0
+    assert data["diff_hints"]
+
+
+def test_analyze_agent_trace_sample_html(tmp_path):
+    sample = Path(__file__).parents[1] / "examples" / "agent-trace" / "sample.json"
+    out = tmp_path / "agent-trace-report.html"
+    runner = CliRunner()
+    result = runner.invoke(main, ["analyze", str(sample), "--format", "agent-trace", "--html", str(out)])
+    assert result.exit_code == 0
+    assert out.exists()
+    assert "<html" in out.read_text().lower()