traceforge-toolkit 0.1.0__py3-none-any.whl

traceforge/preprocessors/codex.py

@@ -0,0 +1,311 @@
+"""Codex CLI preprocessor — normalize JSONL rollout lines into tool call events."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from traceforge.preprocessors.registry import register_preprocessor
+
+
+@register_preprocessor("codex")
+def preprocess_codex(obj: dict[str, Any]) -> list[dict[str, Any]]:
+    """Normalize a Codex CLI rollout line into traceforge-friendly dicts.
+
+    Codex CLI JSONL format (~/.codex/sessions/*.jsonl):
+    Each line: {"timestamp": "...", "type": "...", "payload": {...}}
+
+    Type discriminator values:
+    - "session_meta" — session metadata (first line)
+    - "response_item" — model outputs (function_call, function_call_output, message)
+    - "event_msg" — lifecycle events (exec_command_begin/end, mcp_tool_call_begin/end)
+
+    We flatten the double-type nesting into a single `block_type` field:
+    - response_item.function_call → "tool.shell_call"
+    - response_item.function_call_output → "tool.shell_result"
+    - event_msg.exec_command_begin → "tool.exec_begin"
+    - event_msg.exec_command_end → "tool.exec_end"
+    - event_msg.mcp_tool_call_begin → "tool.mcp_call"
+    - event_msg.mcp_tool_call_end → "tool.mcp_result"
+    - session_meta → "session.meta"
+    - response_item.message (role=assistant) → "message.assistant"
+    - event_msg.user_message → "message.user"
+    """
+    top_type = obj.get("type", "")
+    payload = obj.get("payload") or {}
+    timestamp = obj.get("timestamp", "")
+
+    # If already preprocessed (has block_type but no rollout structure), pass through
+    if "block_type" in obj and not top_type:
+        return [obj]
+
+    if not isinstance(payload, dict):
+        return [obj]
+
+    base = {"_timestamp": timestamp}
+
+    if top_type == "session_meta":
+        return [
+            {
+                **base,
+                "block_type": "session.meta",
+                "session_id": payload.get("id", ""),
+                "model_provider": payload.get("model_provider", ""),
+                "cwd": payload.get("cwd", ""),
+                "cli_version": payload.get("cli_version", ""),
+            }
+        ]
+
+    if top_type == "response_item":
+        return _handle_response_item(payload, base)
+
+    if top_type == "event_msg":
+        return _handle_event_msg(payload, base)
+
+    if top_type == "turn_context":
+        return []
+
+    # Pass through unknown types
+    return [{**base, "block_type": f"raw.{top_type}", **payload}]
+
+
+def _handle_response_item(payload: dict[str, Any], base: dict[str, Any]) -> list[dict[str, Any]]:
+    """Handle response_item lines (model output)."""
+    item_type = payload.get("type", "")
+
+    if item_type == "function_call":
+        # Shell/built-in tool call from model
+        arguments_raw = payload.get("arguments", "{}")
+        # arguments is a JSON-encoded string for function_call
+        try:
+            arguments = (
+                json.loads(arguments_raw) if isinstance(arguments_raw, str) else arguments_raw
+            )
+        except (json.JSONDecodeError, TypeError):
+            arguments = {"_raw": arguments_raw}
+
+        return [
+            {
+                **base,
+                "block_type": "tool.shell_call",
+                "call_id": payload.get("call_id", ""),
+                "tool_name": payload.get("name", "shell"),
+                "arguments": arguments,
+                "namespace": payload.get("namespace"),
+            }
+        ]
+
+    if item_type == "function_call_output":
+        output = payload.get("output", "")
+        return [
+            {
+                **base,
+                "block_type": "tool.shell_result",
+                "call_id": payload.get("call_id", ""),
+                "output": output,
+            }
+        ]
+
+    if item_type == "message":
+        role = payload.get("role", "")
+        content_blocks = payload.get("content", [])
+        if role == "developer":
+            return []
+
+        text = ""
+        if isinstance(content_blocks, list):
+            text = " ".join(
+                b.get("text", "")
+                for b in content_blocks
+                if isinstance(b, dict) and b.get("type") in {"output_text", "input_text"}
+            )
+        elif isinstance(content_blocks, str):
+            text = content_blocks
+        return [
+            {
+                **base,
+                "block_type": f"message.{role}" if role else "message.unknown",
+                "content": text,
+            }
+        ]
+
+    if item_type == "reasoning":
+        return []
+
+    # Other response items (e.g. local_shell_call, custom_tool_call)
+    if item_type == "local_shell_call":
+        action = payload.get("action", {})
+        return [
+            {
+                **base,
+                "block_type": "tool.shell_call",
+                "call_id": payload.get("call_id", ""),
+                "tool_name": "shell",
+                "arguments": action if isinstance(action, dict) else {"_raw": action},
+            }
+        ]
+
+    if item_type == "custom_tool_call":
+        input_raw = payload.get("input", "{}")
+        try:
+            arguments = json.loads(input_raw) if isinstance(input_raw, str) else input_raw
+        except (json.JSONDecodeError, TypeError):
+            arguments = {"_raw": input_raw}
+        return [
+            {
+                **base,
+                "block_type": "tool.custom_call",
+                "call_id": payload.get("call_id", ""),
+                "tool_name": payload.get("name", ""),
+                "arguments": arguments,
+            }
+        ]
+
+    if item_type == "custom_tool_call_output":
+        return [
+            {
+                **base,
+                "block_type": "tool.custom_result",
+                "call_id": payload.get("call_id", ""),
+                "tool_name": payload.get("name", ""),
+                "output": payload.get("output", ""),
+            }
+        ]
+
+    return [
+        {
+            **base,
+            "block_type": f"response.{item_type}",
+            **{k: v for k, v in payload.items() if k != "type"},
+        }
+    ]
+
+
+def _handle_event_msg(payload: dict[str, Any], base: dict[str, Any]) -> list[dict[str, Any]]:
+    """Handle event_msg lines (lifecycle events)."""
+    event_type = payload.get("type", "")
+
+    if event_type == "exec_command_begin":
+        return [
+            {
+                **base,
+                "block_type": "tool.exec_begin",
+                "call_id": payload.get("call_id", ""),
+                "command": payload.get("command", []),
+                "cwd": payload.get("cwd", ""),
+            }
+        ]
+
+    if event_type == "exec_command_end":
+        return [
+            {
+                **base,
+                "block_type": "tool.exec_end",
+                "call_id": payload.get("call_id", ""),
+                "command": payload.get("command", []),
+                "cwd": payload.get("cwd", ""),
+                "exit_code": payload.get("exit_code"),
+                "stdout": payload.get("stdout", ""),
+                "stderr": payload.get("stderr", ""),
+                "status": payload.get("status", ""),
+            }
+        ]
+
+    if event_type == "patch_apply_begin":
+        return [
+            {
+                **base,
+                "block_type": "tool.exec_begin",
+                "call_id": payload.get("call_id", ""),
+                "command": "apply_patch",
+                "cwd": payload.get("cwd", ""),
+            }
+        ]
+
+    if event_type == "patch_apply_end":
+        return [
+            {
+                **base,
+                "block_type": "tool.exec_end",
+                "call_id": payload.get("call_id", ""),
+                "command": "apply_patch",
+                "cwd": payload.get("cwd", ""),
+                "exit_code": 0 if payload.get("success") else 1,
+                "stdout": payload.get("stdout", ""),
+                "stderr": payload.get("stderr", ""),
+                "status": payload.get("status", ""),
+                "changes": payload.get("changes", {}),
+            }
+        ]
+
+    if event_type == "mcp_tool_call_begin":
+        invocation = payload.get("invocation") or {}
+        return [
+            {
+                **base,
+                "block_type": "tool.mcp_call",
+                "call_id": payload.get("call_id", ""),
+                "tool_name": invocation.get("tool", ""),
+                "server": invocation.get("server", ""),
+                "arguments": invocation.get("arguments"),
+            }
+        ]
+
+    if event_type == "mcp_tool_call_end":
+        invocation = payload.get("invocation") or {}
+        result = payload.get("result", {})
+        # Result is {"Ok": {...}} or {"Err": "..."}
+        ok_val = result.get("Ok", {}) if isinstance(result, dict) else {}
+        err_val = result.get("Err") if isinstance(result, dict) else None
+        return [
+            {
+                **base,
+                "block_type": "tool.mcp_result",
+                "call_id": payload.get("call_id", ""),
+                "tool_name": invocation.get("tool", ""),
+                "server": invocation.get("server", ""),
+                "is_error": err_val is not None
+                or (isinstance(ok_val, dict) and ok_val.get("is_error", False)),
+                "output": err_val if err_val else ok_val,
+            }
+        ]
+
+    if event_type == "exec_approval_request":
+        return [
+            {
+                **base,
+                "block_type": "tool.approval_request",
+                "call_id": payload.get("call_id", ""),
+                "command": payload.get("command", ""),
+            }
+        ]
+
+    if event_type == "user_message":
+        return [
+            {
+                **base,
+                "block_type": "message.user",
+                "content": payload.get("message", ""),
+            }
+        ]
+
+    if event_type == "agent_message":
+        return [
+            {
+                **base,
+                "block_type": "message.assistant",
+                "content": payload.get("message", ""),
+            }
+        ]
+
+    if event_type in {"task_started", "task_complete", "token_count"}:
+        return []
+
+    # Pass through other event types
+    return [
+        {
+            **base,
+            "block_type": f"event.{event_type}",
+            **{k: v for k, v in payload.items() if k != "type"},
+        }
+    ]

traceforge/preprocessors/continue_dev.py

@@ -0,0 +1,119 @@
+"""Continue.dev preprocessor — flatten session JSON history into per-message events."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from traceforge.preprocessors.registry import register_preprocessor
+
+
+@register_preprocessor("continue")
+def preprocess_continue(obj: dict[str, Any]) -> list[dict[str, Any]]:
+    """Flatten Continue.dev session JSON into per-message dicts.
+
+    Continue.dev session format (~/.continue/sessions/{id}.json):
+    {
+      "sessionId": "...",
+      "title": "...",
+      "history": [
+        {"message": {"role": "user", "content": "..."}},
+        {"message": {"role": "assistant", "content": "...", "toolCalls": [...]}},
+        {"message": {"role": "tool", "toolCallId": "...", "content": "..."}}
+      ]
+    }
+
+    Note: Continue.dev persists camelCase keys (toolCalls, toolCallId) — its
+    internal TypeScript ChatMessage types are written to disk via JSON.stringify
+    with no rename. snake_case (tool_calls/tool_call_id) only appears on the
+    OpenAI provider wire format, not in the session file on disk.
+
+    Each history entry becomes a normalized dict with:
+    - block_type: "user.message", "assistant.message", "assistant.tool_use", "tool.result"
+    - Flattened message fields
+    """
+    history = obj.get("history")
+    if not isinstance(history, list):
+        return [obj]
+
+    session_id = obj.get("sessionId", "")
+    results: list[dict[str, Any]] = []
+
+    for entry in history:
+        message = entry.get("message") if isinstance(entry, dict) else None
+        if not isinstance(message, dict):
+            continue
+
+        role = message.get("role", "")
+        content = message.get("content")
+        # Continue.dev persists session JSON with camelCase keys (TypeScript
+        # JSON.stringify of its internal ChatMessage types). snake_case keys
+        # only appear on the OpenAI provider wire format, never on disk.
+        tool_calls = message.get("toolCalls")
+
+        if role == "user":
+            results.append(
+                {
+                    "block_type": "user.message",
+                    "session_id": session_id,
+                    "content": content,
+                }
+            )
+
+        elif role == "assistant":
+            # Assistant messages may contain both text and tool_calls
+            if content:
+                results.append(
+                    {
+                        "block_type": "assistant.message",
+                        "session_id": session_id,
+                        "content": content,
+                    }
+                )
+
+            if isinstance(tool_calls, list):
+                for tc in tool_calls:
+                    if not isinstance(tc, dict):
+                        continue
+                    func = tc.get("function", {})
+                    args_raw = func.get("arguments", "{}") if isinstance(func, dict) else "{}"
+                    # Parse arguments string to dict for consistency with other preprocessors
+                    if isinstance(args_raw, str):
+                        try:
+                            arguments = json.loads(args_raw)
+                        except (json.JSONDecodeError, ValueError):
+                            arguments = {"_raw": args_raw}
+                    else:
+                        arguments = args_raw if isinstance(args_raw, dict) else {}
+                    results.append(
+                        {
+                            "block_type": "assistant.tool_use",
+                            "session_id": session_id,
+                            "tool_call_id": tc.get("id", ""),
+                            "tool_name": func.get("name", "") if isinstance(func, dict) else "",
+                            "arguments": arguments,
+                        }
+                    )
+
+        elif role == "thinking":
+            # Claude extended thinking / OpenAI reasoning tokens persisted on disk
+            # as role "thinking" (ThinkingChatMessage). Without this they are dropped.
+            results.append(
+                {
+                    "block_type": "assistant.thinking",
+                    "session_id": session_id,
+                    "content": content,
+                }
+            )
+
+        elif role == "tool":
+            results.append(
+                {
+                    "block_type": "tool.result",
+                    "session_id": session_id,
+                    "tool_call_id": message.get("toolCallId", ""),
+                    "content": content,
+                }
+            )
+
+    return results if results else [obj]

traceforge/preprocessors/copilot_vscode.py

@@ -0,0 +1,171 @@
+"""VS Code Copilot Chat preprocessor — replay the ChatModel journal.
+
+VS Code persists each Copilot Chat session as a line-delimited journal at
+``workspaceStorage/<hash>/chatSessions/<sessionId>.jsonl`` (ChatModel version 3).
+Each physical line is one journal record:
+
+* ``{"kind": 0, "v": {...}}``            — full snapshot (session start)
+* ``{"kind": 1, "k": [path...], "v": x}`` — *set* value ``x`` at ``path``
+* ``{"kind": 2, "k": [path...], "v": [...]}`` — *append* (list-extend) at ``path``
+
+Paths mix dict keys and integer array indices, e.g.
+``["requests", 2, "response"]``. The interesting state lives under ``requests[]``:
+each request carries the user ``message``, a streamed ``response`` part list
+(``thinking`` / ``toolInvocationSerialized`` / markdown text / file refs), and a
+terminal ``result`` with timings.
+
+Because the adapter feeds records one physical line at a time, this preprocessor
+keeps a small amount of module-level state (a mirror of the ``requests`` index ->
+metadata map) so streamed response parts can be attributed to their originating
+request. State is reset whenever a snapshot (``kind == 0``) is seen, which is the
+first record of every session file, so state never bleeds across sessions.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from traceforge.preprocessors.registry import register_preprocessor
+
+# Per-session reconstruction state (reset on every snapshot record).
+_REQ_IDS: dict[int, str | None] = {}
+_REQ_MODELS: dict[int, Any] = {}
+_REQ_TS: dict[int, Any] = {}
+_REQ_COUNT = [0]  # boxed int so helpers can mutate it
+
+
+def _reset() -> None:
+    _REQ_IDS.clear()
+    _REQ_MODELS.clear()
+    _REQ_TS.clear()
+    _REQ_COUNT[0] = 0
+
+
+def _agent_id(req: dict[str, Any]) -> Any:
+    agent = req.get("agent")
+    if isinstance(agent, dict):
+        return agent.get("id") or agent.get("name")
+    return agent
+
+
+def _emit_part(part: Any, idx: int) -> dict[str, Any] | None:
+    """Turn one streamed ``response`` part into a flat, typed dict."""
+    if not isinstance(part, dict):
+        return None
+    flat = dict(part)
+    # Markdown content parts have no ``kind`` discriminator; everything else
+    # (thinking / toolInvocationSerialized / inlineReference / ...) does.
+    flat["event_type"] = part.get("kind") or "assistant_text"
+    flat["request_id"] = _REQ_IDS.get(idx)
+    flat["model"] = _REQ_MODELS.get(idx)
+    flat["timestamp"] = _REQ_TS.get(idx)
+    return flat
+
+
+def _emit_request(req: dict[str, Any], idx: int) -> list[dict[str, Any]]:
+    """Register a request at ``idx`` and emit its user message + any inline parts."""
+    _REQ_IDS[idx] = req.get("requestId")
+    _REQ_MODELS[idx] = req.get("modelId")
+    _REQ_TS[idx] = req.get("timestamp")
+
+    message = req.get("message")
+    text = message.get("text") if isinstance(message, dict) else message
+
+    out: list[dict[str, Any]] = [
+        {
+            "event_type": "user_message",
+            "request_id": req.get("requestId"),
+            "model": req.get("modelId"),
+            "agent": _agent_id(req),
+            "text": text,
+            "timestamp": req.get("timestamp"),
+        }
+    ]
+    # A compacted journal may carry the full response/result inline on the
+    # request itself; a streamed one appends them later. Emit whatever is here.
+    for part in req.get("response") or []:
+        emitted = _emit_part(part, idx)
+        if emitted is not None:
+            out.append(emitted)
+    result = req.get("result")
+    if isinstance(result, dict):
+        out.append(_emit_result(result, idx))
+    return out
+
+
+def _emit_result(result: dict[str, Any], idx: int) -> dict[str, Any]:
+    timings = result.get("timings") if isinstance(result.get("timings"), dict) else {}
+    return {
+        "event_type": "request_result",
+        "request_id": _REQ_IDS.get(idx),
+        "model": _REQ_MODELS.get(idx),
+        "first_progress_ms": timings.get("firstProgress"),
+        "total_elapsed_ms": timings.get("totalElapsed"),
+        "timestamp": _REQ_TS.get(idx),
+    }
+
+
+@register_preprocessor("copilot_vscode")
+def preprocess_copilot_vscode(obj: dict[str, Any]) -> list[dict[str, Any]]:
+    """Expand one VS Code ChatModel journal record into typed flat events."""
+    # Passthrough for already-typed rows (generic conformance probes feed a line
+    # keyed by the post-preprocessor type_field, with no journal ``kind``).
+    if "kind" not in obj and obj.get("event_type"):
+        return [obj]
+
+    kind = obj.get("kind")
+
+    # ── Snapshot: session start (+ any requests already present) ──────────────
+    if kind == 0:
+        _reset()
+        snap = obj.get("v") if isinstance(obj.get("v"), dict) else {}
+        out: list[dict[str, Any]] = [
+            {
+                "event_type": "session_started",
+                "session_id": snap.get("sessionId"),
+                "version": snap.get("version"),
+                "responder": snap.get("responderUsername"),
+                "initial_location": snap.get("initialLocation"),
+                "timestamp": snap.get("creationDate"),
+            }
+        ]
+        for req in snap.get("requests") or []:
+            if isinstance(req, dict):
+                out.extend(_emit_request(req, _REQ_COUNT[0]))
+                _REQ_COUNT[0] += 1
+        return out
+
+    path = obj.get("k")
+    value = obj.get("v")
+    if not isinstance(path, list) or not path:
+        return []
+
+    # ── Append (list-extend) records ──────────────────────────────────────────
+    if kind == 2:
+        # New requests appended to the top-level requests array.
+        if path == ["requests"]:
+            out = []
+            for req in value or []:
+                if isinstance(req, dict):
+                    out.extend(_emit_request(req, _REQ_COUNT[0]))
+                    _REQ_COUNT[0] += 1
+            return out
+        # Streamed response parts appended to requests[i].response.
+        if len(path) == 3 and path[0] == "requests" and path[2] == "response":
+            idx = path[1]
+            out = []
+            for part in value or []:
+                emitted = _emit_part(part, idx)
+                if emitted is not None:
+                    out.append(emitted)
+            return out
+        return []
+
+    # ── Set records: only the terminal per-request result is interesting ──────
+    if kind == 1:
+        if len(path) == 3 and path[0] == "requests" and path[2] == "result":
+            if isinstance(value, dict):
+                return [_emit_result(value, path[1])]
+        return []
+
+    return []