From 5c3b1539ee231a08298e6ba0440a0fb014621994 Mon Sep 17 00:00:00 2001 From: Varshith Date: Wed, 3 Jun 2026 20:15:29 -0500 Subject: [PATCH 1/2] feat(annotate-mlflow-trace): add trace tagging and feedback skill Signed-off-by: Varshith --- .claude-plugin/plugin.json | 1 + README.md | 1 + annotate-mlflow-trace/SKILL.md | 72 ++++++ .../scripts/test_trace_annotate.py | 106 +++++++++ .../scripts/trace_annotate.py | 224 ++++++++++++++++++ hooks/README.md | 27 +++ hooks/mlflow-suggest-hook.py | 3 + hooks/skip_skill_traces.py | 96 ++++++++ hooks/test_skip_skill_traces.py | 105 ++++++++ 9 files changed, 635 insertions(+) create mode 100644 annotate-mlflow-trace/SKILL.md create mode 100644 annotate-mlflow-trace/scripts/test_trace_annotate.py create mode 100644 annotate-mlflow-trace/scripts/trace_annotate.py create mode 100644 hooks/skip_skill_traces.py create mode 100644 hooks/test_skip_skill_traces.py diff --git a/.claude-plugin/plugin.json b/.claude-plugin/plugin.json index c4d611e..9cae7d7 100644 --- a/.claude-plugin/plugin.json +++ b/.claude-plugin/plugin.json @@ -12,6 +12,7 @@ "analyze-mlflow-trace", "analyze-mlflow-chat-session", "retrieving-mlflow-traces", + "annotate-mlflow-trace", "querying-mlflow-metrics", "mlflow-onboarding", "searching-mlflow-docs" diff --git a/README.md b/README.md index c576485..b6ebae9 100644 --- a/README.md +++ b/README.md @@ -30,6 +30,7 @@ Building production-ready AI agents is hard. You need observability to understan | **analyze-mlflow-trace** | Debugs issues by examining spans, assessments, and correlating with your codebase. | | **analyze-mlflow-chat-session** | Debugs multi-turn chat conversations by reconstructing session history and finding where things went wrong. | | **retrieving-mlflow-traces** | Powerful trace search and filtering by status, session, user, time range, or custom metadata. | +| **annotate-mlflow-trace** | Tags and logs human feedback on the current session's trace inline, so coding sessions can be labeled for later evaluation. | ### Evaluation & Metrics diff --git a/annotate-mlflow-trace/SKILL.md b/annotate-mlflow-trace/SKILL.md new file mode 100644 index 0000000..6ecd704 --- /dev/null +++ b/annotate-mlflow-trace/SKILL.md @@ -0,0 +1,72 @@ +--- +name: annotate-mlflow-trace +description: Tags and logs human feedback on the current coding session's MLflow trace. Triggers on requests to tag a trace, label a trace, rate a trace, leave feedback, thumbs up/down, or mark a session for later evaluation. +--- + +# Annotate MLflow Trace + +Annotate the current coding session's MLflow trace inline, so sessions can be labeled and rated +for later evaluation without switching to the MLflow UI. Run `scripts/trace_annotate.py` with the +`tag`, `feedback`, or `list` subcommand. + +Tracking config is read from the environment (`MLFLOW_TRACKING_URI`, `MLFLOW_EXPERIMENT_ID`, +`MLFLOW_EXPERIMENT_NAME`). Every subcommand defaults to the most recent trace; pass `--trace-id` +to target a specific one. Run `list` first if you need to find the trace. + +## Tag a trace + +```bash +# Tag the most recent trace +uv run python ${CLAUDE_SKILL_DIR}/scripts/trace_annotate.py tag quality=good + +# Multiple tags at once +uv run python ${CLAUDE_SKILL_DIR}/scripts/trace_annotate.py tag sprint=42 reviewer=alice bug=true + +# Tag a specific trace by ID +uv run python ${CLAUDE_SKILL_DIR}/scripts/trace_annotate.py tag --trace-id tr-abc123 quality=good +``` + +## Leave feedback + +Logs a human assessment via `mlflow.log_feedback()` with a `HUMAN` source and sets a +`has_feedback=true` tag so feedback-bearing traces are easy to search for. + +```bash +# Thumbs up on the current session trace +uv run python ${CLAUDE_SKILL_DIR}/scripts/trace_annotate.py feedback --name thumbs_up --value true + +# Rating with a rationale +uv run python ${CLAUDE_SKILL_DIR}/scripts/trace_annotate.py feedback --name quality --value good \ + --rationale "Completed the refactor correctly" +``` + +Boolean and numeric values are stored as their typed value (`true` becomes a boolean, `0.9` a +float); anything else is stored as text. + +## List recent traces + +```bash +uv run python ${CLAUDE_SKILL_DIR}/scripts/trace_annotate.py list +uv run python ${CLAUDE_SKILL_DIR}/scripts/trace_annotate.py list --limit 20 +``` + +Shows trace ID, time, status, and an input preview so you can identify the trace to annotate. + +## Arguments + +| Subcommand | Argument | Required | Description | +|------------|----------|----------|-------------| +| `tag` | `pairs` | Yes | One or more `key=value` tags | +| `tag` | `--trace-id` | No | Trace to tag (default: most recent) | +| `feedback` | `--name` | Yes | Assessment name, e.g. `thumbs_up`, `quality` | +| `feedback` | `--value` | Yes | Assessment value (`true`/`false`, a number, or text) | +| `feedback` | `--rationale` | No | Justification for the feedback | +| `feedback` | `--source-id` | No | Identifier for the human reviewer | +| `feedback` | `--trace-id` | No | Trace to annotate (default: most recent) | +| `list` | `--limit` | No | Number of traces to show (default: 10) | + +## Keep annotations out of your traces + +The script disables tracing for its own MLflow calls so annotating does not create extra spans. +To also stop the Stop hook from logging a session trace for the annotation turn itself, install +`hooks/skip_skill_traces.py`. See [`hooks/README.md`](../hooks/README.md) for setup. diff --git a/annotate-mlflow-trace/scripts/test_trace_annotate.py b/annotate-mlflow-trace/scripts/test_trace_annotate.py new file mode 100644 index 0000000..1965460 --- /dev/null +++ b/annotate-mlflow-trace/scripts/test_trace_annotate.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python3 +"""Tests for trace_annotate.py. + +Covers the pure parsing and trace-selection logic plus the CLI's argument handling, without +needing a live MLflow tracking store (MLflow is imported lazily, only by the subcommands that +talk to the store). Run with pytest, or standalone: +`python annotate-mlflow-trace/scripts/test_trace_annotate.py`. +""" +from __future__ import annotations + +import subprocess +import sys +import types +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent)) + +import trace_annotate + +SCRIPT = Path(__file__).parent / "trace_annotate.py" + + +def _trace(trace_id, name="claude_code_conversation", request_time_ms=0): + return types.SimpleNamespace( + info=types.SimpleNamespace( + trace_id=trace_id, + tags={trace_annotate.TRACE_NAME_TAG: name}, + request_time=request_time_ms, + ) + ) + + +def test_parse_tag_pairs_splits_key_value(): + assert trace_annotate.parse_tag_pairs(["quality=good", "sprint=42"]) == { + "quality": "good", + "sprint": "42", + } + + +def test_parse_tag_pairs_keeps_value_containing_equals(): + assert trace_annotate.parse_tag_pairs(["note=a=b"]) == {"note": "a=b"} + + +def test_parse_tag_pairs_rejects_missing_equals(): + try: + trace_annotate.parse_tag_pairs(["bad"]) + except ValueError as e: + assert "key=value" in str(e) + else: + raise AssertionError("expected ValueError for a pair without '='") + + +def test_coerce_value_handles_booleans_numbers_and_text(): + assert trace_annotate.coerce_value("true") is True + assert trace_annotate.coerce_value("False") is False + assert trace_annotate.coerce_value("42") == 42 + assert trace_annotate.coerce_value("0.9") == 0.9 + assert trace_annotate.coerce_value("good") == "good" + + +def test_select_most_recent_picks_latest(): + traces = [_trace("tr-old", request_time_ms=100), _trace("tr-new", request_time_ms=200)] + assert trace_annotate.select_most_recent(traces).info.trace_id == "tr-new" + + +def test_select_most_recent_skips_companion_traces(): + traces = [ + _trace("tr-session", request_time_ms=100), + _trace("tr-env", name="env_snapshot", request_time_ms=200), + ] + assert trace_annotate.select_most_recent(traces).info.trace_id == "tr-session" + + +def test_select_most_recent_returns_none_when_empty(): + assert trace_annotate.select_most_recent([]) is None + + +def test_cli_tag_without_pairs_errors(): + proc = subprocess.run([sys.executable, str(SCRIPT), "tag"], capture_output=True, text=True) + assert proc.returncode != 0 + assert "usage" in proc.stderr.lower() + + +def test_cli_help_lists_subcommands(): + proc = subprocess.run([sys.executable, str(SCRIPT), "--help"], capture_output=True, text=True) + assert proc.returncode == 0 + for subcommand in ("tag", "feedback", "list"): + assert subcommand in proc.stdout + + +def _main() -> int: + tests = [v for k, v in sorted(globals().items()) if k.startswith("test_") and callable(v)] + failures = 0 + for test in tests: + try: + test() + print(f"PASS {test.__name__}") + except AssertionError as e: + failures += 1 + print(f"FAIL {test.__name__}: {e}") + print(f"\n{len(tests) - failures}/{len(tests)} passed") + return 1 if failures else 0 + + +if __name__ == "__main__": + raise SystemExit(_main()) diff --git a/annotate-mlflow-trace/scripts/trace_annotate.py b/annotate-mlflow-trace/scripts/trace_annotate.py new file mode 100644 index 0000000..c68dcaa --- /dev/null +++ b/annotate-mlflow-trace/scripts/trace_annotate.py @@ -0,0 +1,224 @@ +#!/usr/bin/env python3 +"""Tag and log human feedback on MLflow traces from the current coding session.""" + +from __future__ import annotations + +import argparse +import sys +from datetime import datetime, timezone + +TRACE_NAME_TAG = "mlflow.traceName" +HAS_FEEDBACK_TAG = "has_feedback" + +# Companion traces logged alongside a session that should not be treated as "the" session +# trace when resolving the most recent one. +COMPANION_TRACE_NAMES = {"env_snapshot"} + + +def parse_tag_pairs(pairs: list[str]) -> dict[str, str]: + """Parse ``key=value`` strings into a tag dictionary.""" + tags = {} + for pair in pairs: + if "=" not in pair: + raise ValueError(f"Invalid tag '{pair}'. Use key=value format, e.g. quality=good.") + key, value = pair.split("=", 1) + key = key.strip() + if not key: + raise ValueError(f"Invalid tag '{pair}'. Tag key cannot be empty.") + tags[key] = value.strip() + return tags + + +def coerce_value(raw: str) -> bool | int | float | str: + """Coerce a feedback value to bool, int, or float when it unambiguously is one.""" + if raw.lower() in ("true", "false"): + return raw.lower() == "true" + try: + return int(raw) + except ValueError: + pass + try: + return float(raw) + except ValueError: + pass + return raw + + +def trace_name(trace) -> str: + """Return the human-readable name of a trace, or an empty string.""" + return (trace.info.tags or {}).get(TRACE_NAME_TAG, "") + + +def _request_time_ms(trace) -> float: + request_time = getattr(trace.info, "request_time", None) + if isinstance(request_time, (int, float)): + return float(request_time) + if isinstance(request_time, datetime): + return request_time.timestamp() * 1000 + return 0.0 + + +def select_most_recent(traces: list, companion_names: set[str] = COMPANION_TRACE_NAMES): + """Return the most recent trace that is not a companion trace, or None.""" + candidates = [t for t in traces if trace_name(t) not in companion_names] + if not candidates: + return None + return max(candidates, key=_request_time_ms) + + +def _format_time(trace) -> str: + request_time = getattr(trace.info, "request_time", None) + if isinstance(request_time, (int, float)): + return datetime.fromtimestamp(request_time / 1000, tz=timezone.utc).strftime( + "%Y-%m-%d %H:%M:%S" + ) + if isinstance(request_time, datetime): + return request_time.astimezone(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + return "unknown" + + +def _format_state(trace) -> str: + state = getattr(trace.info, "state", None) + if state is None: + return "unknown" + return getattr(state, "name", None) or str(state) + + +def _format_preview(trace, width: int = 60) -> str: + preview = " ".join((getattr(trace.info, "request_preview", None) or "").split()) + return preview[: width - 3] + "..." if len(preview) > width else preview + + +def _format_table(traces: list) -> str: + headers = ["TRACE ID", "TIME", "STATUS", "PREVIEW"] + rows = [ + [t.info.trace_id, _format_time(t), _format_state(t), _format_preview(t)] for t in traces + ] + widths = [max(len(h), max((len(r[i]) for r in rows), default=0)) for i, h in enumerate(headers)] + lines = [ + " ".join(h.ljust(widths[i]) for i, h in enumerate(headers)), + " ".join("-" * w for w in widths), + ] + lines.extend(" ".join(cell.ljust(widths[i]) for i, cell in enumerate(row)) for row in rows) + return "\n".join(lines) + + +def _configure_mlflow() -> None: + """Point MLflow at the tracking store and experiment from the environment.""" + import mlflow + from mlflow.claude_code.config import get_env_var + from mlflow.environment_variables import ( + MLFLOW_EXPERIMENT_ID, + MLFLOW_EXPERIMENT_NAME, + MLFLOW_TRACKING_URI, + ) + + # Don't let our own annotation calls generate traces. + mlflow.tracing.disable() + + tracking_uri = get_env_var(MLFLOW_TRACKING_URI.name) + if tracking_uri: + mlflow.set_tracking_uri(tracking_uri) + + experiment_id = get_env_var(MLFLOW_EXPERIMENT_ID.name) + experiment_name = get_env_var(MLFLOW_EXPERIMENT_NAME.name) + if experiment_id: + mlflow.set_experiment(experiment_id=experiment_id) + elif experiment_name: + mlflow.set_experiment(experiment_name) + + +def _resolve_trace_id(args) -> str: + if args.trace_id: + return args.trace_id + + import mlflow + + traces = mlflow.search_traces(max_results=50, return_type="list") + trace = select_most_recent(traces) + if trace is None: + raise RuntimeError("No traces found to annotate. Run a session first, or pass --trace-id.") + return trace.info.trace_id + + +def cmd_tag(args) -> None: + import mlflow + + _configure_mlflow() + tags = parse_tag_pairs(args.pairs) + trace_id = _resolve_trace_id(args) + for key, value in tags.items(): + mlflow.set_trace_tag(trace_id, key, value) + summary = ", ".join(f"{k}={v}" for k, v in tags.items()) + print(f"Tagged {trace_id}: {summary}") + + +def cmd_feedback(args) -> None: + import mlflow + from mlflow.entities import AssessmentSource + + _configure_mlflow() + trace_id = _resolve_trace_id(args) + source = ( + AssessmentSource(source_type="HUMAN", source_id=args.source_id) + if args.source_id + else AssessmentSource(source_type="HUMAN") + ) + mlflow.log_feedback( + trace_id=trace_id, + name=args.name, + value=coerce_value(args.value), + rationale=args.rationale, + source=source, + ) + mlflow.set_trace_tag(trace_id, HAS_FEEDBACK_TAG, "true") + print(f"Logged feedback '{args.name}' on {trace_id}") + + +def cmd_list(args) -> None: + import mlflow + + _configure_mlflow() + traces = mlflow.search_traces(max_results=args.limit, return_type="list") + traces = [t for t in traces if trace_name(t) not in COMPANION_TRACE_NAMES] + if not traces: + print("No traces found.") + return + traces.sort(key=_request_time_ms, reverse=True) + print(_format_table(traces)) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Annotate MLflow traces with tags and feedback") + subparsers = parser.add_subparsers(dest="command", required=True) + + tag_parser = subparsers.add_parser("tag", help="Add key=value tags to a trace") + tag_parser.add_argument("pairs", nargs="+", help="One or more key=value tags") + tag_parser.add_argument("--trace-id", help="Trace to tag (default: most recent)") + tag_parser.set_defaults(func=cmd_tag) + + feedback_parser = subparsers.add_parser("feedback", help="Log human feedback on a trace") + feedback_parser.add_argument("--name", required=True, help="Assessment name, e.g. thumbs_up") + feedback_parser.add_argument( + "--value", required=True, help="Assessment value (true/false, number, or text)" + ) + feedback_parser.add_argument("--rationale", help="Justification for the feedback") + feedback_parser.add_argument("--source-id", help="Identifier for the human reviewer") + feedback_parser.add_argument("--trace-id", help="Trace to annotate (default: most recent)") + feedback_parser.set_defaults(func=cmd_feedback) + + list_parser = subparsers.add_parser("list", help="List recent traces") + list_parser.add_argument("--limit", type=int, default=10, help="Number of traces to show") + list_parser.set_defaults(func=cmd_list) + + args = parser.parse_args() + + try: + args.func(args) + except (ValueError, RuntimeError) as e: + print(f"Error: {e}", file=sys.stderr) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/hooks/README.md b/hooks/README.md index 0c1de53..1ef4e47 100644 --- a/hooks/README.md +++ b/hooks/README.md @@ -15,6 +15,7 @@ A `UserPromptSubmit` hook that detects MLflow-related patterns in your prompts a | trace id, debug trace, why did, what went wrong, analyze trace | `analyze-mlflow-trace` | | session, conversation, chat history, multi-turn | `analyze-mlflow-chat-session` | | search traces, find traces, filter traces, get trace | `retrieving-mlflow-traces` | +| tag trace, annotate trace, label trace, feedback on trace, rate trace, thumbs up, thumbs down | `annotate-mlflow-trace` | | metrics, token usage, latency, cost, usage trend | `querying-mlflow-metrics` | | get started, set up mlflow, onboard, quickstart | `mlflow-onboarding` | | mlflow docs, mlflow api, how to use mlflow | `searching-mlflow-docs` | @@ -68,3 +69,29 @@ The hook exits cleanly with no output when no MLflow patterns are detected, so i ### Compatibility This hook works with any coding agent that supports the `UserPromptSubmit` hook protocol, including Claude Code. Other agents may use different hook registration mechanisms — check your agent's documentation for the equivalent setting. + +## `skip_skill_traces.py` + +A `Stop` hook pre-filter for the `annotate-mlflow-trace` skill. It wraps MLflow's Claude Code Stop hook so that invoking a trace-annotation skill does not log its own session trace. Without it, every `tag` or `feedback` action would create an administrative trace and pollute the experiment. + +### Installation + +This hook replaces the command of your existing MLflow `Stop` hook. If you enabled MLflow Claude Code tracing, you already have a `Stop` hook that runs `stop_hook_handler()`; wrap it by passing it after `--`. + +```json +{ + "hooks": { + "Stop": [ + { + "command": "uv run python \"$CLAUDE_PROJECT_DIR\"/.claude/hooks/skip_skill_traces.py -- uv run python -c \"from mlflow.claude_code.hooks import stop_hook_handler; stop_hook_handler()\"" + } + ] + } +} +``` + +Copy `skip_skill_traces.py` to `.claude/hooks/` (or reference it by absolute path). + +### How It Works + +The hook reads the Stop event from stdin, opens the session transcript, and checks whether the last user turn invoked a skill in its skip list (detected by the `Base directory for this skill:` line Claude Code injects). If so, it short-circuits with `{"continue": true, "suppressOutput": true}` and the wrapped command never runs, so no trace is logged. For every other turn it passes stdin through to the wrapped command unchanged, so normal tracing is unaffected. diff --git a/hooks/mlflow-suggest-hook.py b/hooks/mlflow-suggest-hook.py index 8d16cdb..fcd95ea 100755 --- a/hooks/mlflow-suggest-hook.py +++ b/hooks/mlflow-suggest-hook.py @@ -28,6 +28,9 @@ def main(): if any(k in prompt for k in ["search traces", "find traces", "filter traces", "get trace"]): suggestions.append("💡 Use the `retrieving-mlflow-traces` skill to search/filter traces.") + if any(k in prompt for k in ["tag trace", "annotate trace", "label trace", "feedback on trace", "rate trace", "thumbs up", "thumbs down"]): + suggestions.append("💡 Use the `annotate-mlflow-trace` skill to tag and add feedback to traces.") + if any(k in prompt for k in ["metrics", "token usage", "latency", "cost", "usage trend"]): suggestions.append("💡 Use the `querying-mlflow-metrics` skill to fetch aggregated metrics.") diff --git a/hooks/skip_skill_traces.py b/hooks/skip_skill_traces.py new file mode 100644 index 0000000..ed09554 --- /dev/null +++ b/hooks/skip_skill_traces.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +"""Stop-hook pre-filter that skips trace creation for annotation-skill turns. + +Wraps MLflow's Claude Code Stop hook. When the turn being closed is an invocation of the +`annotate-mlflow-trace` skill, this short-circuits so the annotation action does not log its own +session trace and pollute the experiment with administrative noise (a "thumbs up" should not +create a trace). Every other turn is delegated unchanged to the wrapped command passed after `--`. + +Configure it in `.claude/settings.json`: + + "Stop": [ + { + "command": "uv run python \"$CLAUDE_PROJECT_DIR\"/.claude/hooks/skip_skill_traces.py -- uv run python -c \"from mlflow.claude_code.hooks import stop_hook_handler; stop_hook_handler()\"" + } + ] + +Run the tests with pytest, or standalone: `python hooks/test_skip_skill_traces.py`. +""" +from __future__ import annotations + +import json +import subprocess +import sys + +# Skills whose invocations must not generate their own session trace. +SKIP_SKILLS = {"annotate-mlflow-trace"} + +# Marker Claude Code injects into the user turn when a skill runs. +SKILL_MARKER = "Base directory for this skill:" + + +def _read_transcript(transcript_path: str) -> list[dict]: + entries = [] + with open(transcript_path, encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + entries.append(json.loads(line)) + except json.JSONDecodeError: + continue + return entries + + +def _message_text(entry: dict) -> str: + content = entry.get("message", {}).get("content", "") + if isinstance(content, str): + return content + if isinstance(content, list): + parts = [b.get("text", "") if isinstance(b, dict) else b for b in content] + return "\n".join(p for p in parts if isinstance(p, str)) + return "" + + +def invoked_skill(transcript_path: str) -> str | None: + """Return the name of the skill invoked in the last user turn, or None.""" + last_user_text = "" + for entry in _read_transcript(transcript_path): + if entry.get("type") == "user": + last_user_text = _message_text(entry) + for line in last_user_text.splitlines(): + if SKILL_MARKER in line: + path = line.split(SKILL_MARKER, 1)[1].strip().rstrip("/\\") + return path.replace("\\", "/").split("/")[-1] + return None + + +def main() -> int: + argv = sys.argv[1:] + delegate = argv[argv.index("--") + 1 :] if "--" in argv else [] + + raw = sys.stdin.read() + try: + transcript_path = json.loads(raw).get("transcript_path") + except (json.JSONDecodeError, ValueError): + transcript_path = None + + if transcript_path: + try: + skill = invoked_skill(transcript_path) + except OSError: + skill = None + if skill in SKIP_SKILLS: + print(json.dumps({"continue": True, "suppressOutput": True})) + return 0 + + if not delegate: + print(json.dumps({"continue": True})) + return 0 + + return subprocess.run(delegate, input=raw, text=True).returncode + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/hooks/test_skip_skill_traces.py b/hooks/test_skip_skill_traces.py new file mode 100644 index 0000000..7148fa4 --- /dev/null +++ b/hooks/test_skip_skill_traces.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +"""Tests for skip_skill_traces.py. + +Covers skill detection from a transcript and the stdin/stdout delegation contract: an +annotation-skill turn is short-circuited (the wrapped command never runs), and every other turn +is delegated. Run with pytest, or standalone: `python hooks/test_skip_skill_traces.py`. +""" +from __future__ import annotations + +import contextlib +import json +import subprocess +import sys +import tempfile +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent)) + +import skip_skill_traces + +HOOK = Path(__file__).parent / "skip_skill_traces.py" + + +def _user_line(text): + return json.dumps({"type": "user", "message": {"role": "user", "content": text}}) + + +@contextlib.contextmanager +def transcript(*entries): + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "transcript.jsonl" + path.write_text("\n".join(entries) + "\n", encoding="utf-8") + yield Path(path) + + +def _writer_delegate(target: Path): + return [sys.executable, "-c", f"open({json.dumps(target.as_posix())}, 'w').close()"] + + +def _run(event, delegate=None): + cmd = [sys.executable, str(HOOK)] + if delegate is not None: + cmd += ["--", *delegate] + return subprocess.run(cmd, input=json.dumps(event), capture_output=True, text=True) + + +def test_detects_annotation_skill_invocation(): + marker = f"{skip_skill_traces.SKILL_MARKER} /repo/annotate-mlflow-trace" + with transcript(_user_line("earlier turn"), _user_line(marker)) as path: + assert skip_skill_traces.invoked_skill(str(path)) == "annotate-mlflow-trace" + + +def test_normal_turn_has_no_skill(): + with transcript(_user_line("fix the bug in foo.py")) as path: + assert skip_skill_traces.invoked_skill(str(path)) is None + + +def test_detects_skill_from_block_content(): + blocks = [{"type": "text", "text": f"{skip_skill_traces.SKILL_MARKER} /x/tag-trace"}] + line = json.dumps({"type": "user", "message": {"content": blocks}}) + with transcript(line) as path: + assert skip_skill_traces.invoked_skill(str(path)) == "tag-trace" + + +def test_annotation_turn_short_circuits_without_delegating(): + marker = f"{skip_skill_traces.SKILL_MARKER} /repo/annotate-mlflow-trace" + with transcript(_user_line(marker)) as path: + sentinel = path.parent / "delegated" + proc = _run({"transcript_path": str(path)}, _writer_delegate(sentinel)) + assert proc.returncode == 0, proc.stderr + assert json.loads(proc.stdout)["suppressOutput"] is True + assert not sentinel.exists() + + +def test_normal_turn_delegates(): + with transcript(_user_line("just a normal prompt")) as path: + sentinel = path.parent / "delegated" + proc = _run({"transcript_path": str(path)}, _writer_delegate(sentinel)) + assert proc.returncode == 0, proc.stderr + assert sentinel.exists() + + +def test_malformed_stdin_does_not_crash(): + proc = subprocess.run( + [sys.executable, str(HOOK)], input="not json", capture_output=True, text=True + ) + assert proc.returncode == 0 + + +def _main() -> int: + tests = [v for k, v in sorted(globals().items()) if k.startswith("test_") and callable(v)] + failures = 0 + for test in tests: + try: + test() + print(f"PASS {test.__name__}") + except AssertionError as e: + failures += 1 + print(f"FAIL {test.__name__}: {e}") + print(f"\n{len(tests) - failures}/{len(tests)} passed") + return 1 if failures else 0 + + +if __name__ == "__main__": + raise SystemExit(_main()) From fb9ebcb8a0907557605a71a43563760675d746dd Mon Sep 17 00:00:00 2001 From: Varshith Date: Fri, 26 Jun 2026 19:13:07 -0500 Subject: [PATCH 2/2] feat(annotate-mlflow-trace): scope trace selection to the current project Signed-off-by: Varshith --- annotate-mlflow-trace/SKILL.md | 20 +++-- .../scripts/test_trace_annotate.py | 41 ++++++++-- .../scripts/trace_annotate.py | 76 ++++++++++++++++--- 3 files changed, 114 insertions(+), 23 deletions(-) diff --git a/annotate-mlflow-trace/SKILL.md b/annotate-mlflow-trace/SKILL.md index 6ecd704..19f251f 100644 --- a/annotate-mlflow-trace/SKILL.md +++ b/annotate-mlflow-trace/SKILL.md @@ -10,8 +10,9 @@ for later evaluation without switching to the MLflow UI. Run `scripts/trace_anno `tag`, `feedback`, or `list` subcommand. Tracking config is read from the environment (`MLFLOW_TRACKING_URI`, `MLFLOW_EXPERIMENT_ID`, -`MLFLOW_EXPERIMENT_NAME`). Every subcommand defaults to the most recent trace; pass `--trace-id` -to target a specific one. Run `list` first if you need to find the trace. +`MLFLOW_EXPERIMENT_NAME`). Every subcommand defaults to the most recent trace from the current +project, so a shared experiment does not surface another project's trace. Pass `--trace-id` to +target a specific one, or run `list` first to find it. ## Tag a trace @@ -50,19 +51,20 @@ uv run python ${CLAUDE_SKILL_DIR}/scripts/trace_annotate.py list uv run python ${CLAUDE_SKILL_DIR}/scripts/trace_annotate.py list --limit 20 ``` -Shows trace ID, time, status, and an input preview so you can identify the trace to annotate. +Shows trace ID, time, status, and an input preview for the current project's recent sessions, so +you can identify the trace to annotate. ## Arguments | Subcommand | Argument | Required | Description | |------------|----------|----------|-------------| | `tag` | `pairs` | Yes | One or more `key=value` tags | -| `tag` | `--trace-id` | No | Trace to tag (default: most recent) | +| `tag` | `--trace-id` | No | Trace to tag (default: current project's most recent) | | `feedback` | `--name` | Yes | Assessment name, e.g. `thumbs_up`, `quality` | | `feedback` | `--value` | Yes | Assessment value (`true`/`false`, a number, or text) | | `feedback` | `--rationale` | No | Justification for the feedback | | `feedback` | `--source-id` | No | Identifier for the human reviewer | -| `feedback` | `--trace-id` | No | Trace to annotate (default: most recent) | +| `feedback` | `--trace-id` | No | Trace to annotate (default: current project's most recent) | | `list` | `--limit` | No | Number of traces to show (default: 10) | ## Keep annotations out of your traces @@ -70,3 +72,11 @@ Shows trace ID, time, status, and an input preview so you can identify the trace The script disables tracing for its own MLflow calls so annotating does not create extra spans. To also stop the Stop hook from logging a session trace for the annotation turn itself, install `hooks/skip_skill_traces.py`. See [`hooks/README.md`](../hooks/README.md) for setup. + +## Requirements + +Targets Claude Code sessions traced by `mlflow.claude_code`, and uses the `mlflow.log_feedback` +assessments API. Verified on MLflow 3.12 and 3.14. Project-scoped trace selection reads the +`mlflow.trace.working_directory` metadata written in MLflow 3.14+ and falls back to the most +recent trace on older versions. To annotate a trace from another agent or project, pass +`--trace-id`. diff --git a/annotate-mlflow-trace/scripts/test_trace_annotate.py b/annotate-mlflow-trace/scripts/test_trace_annotate.py index 1965460..38ade4e 100644 --- a/annotate-mlflow-trace/scripts/test_trace_annotate.py +++ b/annotate-mlflow-trace/scripts/test_trace_annotate.py @@ -20,11 +20,12 @@ SCRIPT = Path(__file__).parent / "trace_annotate.py" -def _trace(trace_id, name="claude_code_conversation", request_time_ms=0): +def _trace(trace_id, name="claude_code_conversation", request_time_ms=0, working_directory=""): return types.SimpleNamespace( info=types.SimpleNamespace( trace_id=trace_id, tags={trace_annotate.TRACE_NAME_TAG: name}, + trace_metadata={trace_annotate.WORKING_DIRECTORY_METADATA: working_directory}, request_time=request_time_ms, ) ) @@ -58,21 +59,47 @@ def test_coerce_value_handles_booleans_numbers_and_text(): assert trace_annotate.coerce_value("good") == "good" -def test_select_most_recent_picks_latest(): +def test_select_session_trace_picks_latest(): traces = [_trace("tr-old", request_time_ms=100), _trace("tr-new", request_time_ms=200)] - assert trace_annotate.select_most_recent(traces).info.trace_id == "tr-new" + assert trace_annotate.select_session_trace(traces).info.trace_id == "tr-new" -def test_select_most_recent_skips_companion_traces(): +def test_select_session_trace_skips_companion_traces(): traces = [ _trace("tr-session", request_time_ms=100), _trace("tr-env", name="env_snapshot", request_time_ms=200), ] - assert trace_annotate.select_most_recent(traces).info.trace_id == "tr-session" + assert trace_annotate.select_session_trace(traces).info.trace_id == "tr-session" -def test_select_most_recent_returns_none_when_empty(): - assert trace_annotate.select_most_recent([]) is None +def test_select_session_trace_returns_none_when_empty(): + assert trace_annotate.select_session_trace([]) is None + + +def test_select_session_trace_prefers_current_project(): + traces = [ + _trace("tr-here", request_time_ms=100, working_directory="/projects/alpha"), + _trace("tr-elsewhere", request_time_ms=200, working_directory="/projects/beta"), + ] + selected = trace_annotate.select_session_trace(traces, working_directory="/projects/alpha") + assert selected.info.trace_id == "tr-here" + + +def test_select_session_trace_falls_back_when_no_project_match(): + traces = [ + _trace("tr-old", request_time_ms=100, working_directory="/projects/beta"), + _trace("tr-new", request_time_ms=200, working_directory="/projects/gamma"), + ] + selected = trace_annotate.select_session_trace(traces, working_directory="/projects/alpha") + assert selected.info.trace_id == "tr-new" + + +def test_select_session_trace_prefers_conversation_trace_by_name(): + traces = [ + _trace("tr-conv", name="claude_code_conversation", request_time_ms=100), + _trace("tr-other", name="some_other_trace", request_time_ms=200), + ] + assert trace_annotate.select_session_trace(traces).info.trace_id == "tr-conv" def test_cli_tag_without_pairs_errors(): diff --git a/annotate-mlflow-trace/scripts/trace_annotate.py b/annotate-mlflow-trace/scripts/trace_annotate.py index c68dcaa..a81b08b 100644 --- a/annotate-mlflow-trace/scripts/trace_annotate.py +++ b/annotate-mlflow-trace/scripts/trace_annotate.py @@ -6,11 +6,19 @@ import argparse import sys from datetime import datetime, timezone +from pathlib import Path TRACE_NAME_TAG = "mlflow.traceName" HAS_FEEDBACK_TAG = "has_feedback" +WORKING_DIRECTORY_METADATA = "mlflow.trace.working_directory" -# Companion traces logged alongside a session that should not be treated as "the" session +# Name of the root trace MLflow logs for a Claude Code session. +SESSION_TRACE_NAME = "claude_code_conversation" + +# How many recent traces to scan when resolving the current session's trace. +SEARCH_WINDOW = 50 + +# Companion traces logged alongside a session that should not be treated as the session # trace when resolving the most recent one. COMPANION_TRACE_NAMES = {"env_snapshot"} @@ -49,6 +57,22 @@ def trace_name(trace) -> str: return (trace.info.tags or {}).get(TRACE_NAME_TAG, "") +def trace_working_directory(trace) -> str: + """Return the working directory the trace was created in, or an empty string.""" + metadata = getattr(trace.info, "trace_metadata", None) or {} + return metadata.get(WORKING_DIRECTORY_METADATA, "") + + +def _same_directory(path: str, other: str) -> bool: + """Return True if both strings name the same directory. + + Tolerant of separators and, on Windows, case. False if either side is empty. + """ + if not path or not other: + return False + return Path(path) == Path(other) + + def _request_time_ms(trace) -> float: request_time = getattr(trace.info, "request_time", None) if isinstance(request_time, (int, float)): @@ -58,9 +82,34 @@ def _request_time_ms(trace) -> float: return 0.0 -def select_most_recent(traces: list, companion_names: set[str] = COMPANION_TRACE_NAMES): - """Return the most recent trace that is not a companion trace, or None.""" +def session_candidates( + traces: list, + working_directory: str | None = None, + companion_names: set[str] = COMPANION_TRACE_NAMES, +) -> list: + """Narrow traces to the current coding session's candidates. + + Drops companion traces (e.g. ``env_snapshot``). When ``working_directory`` is given and + any trace was created there, keeps only those, so a shared experiment does not surface + another project's traces. Then prefers the ``claude_code_conversation`` root trace by + name. Each narrowing is skipped when it would drop every candidate, so the result + degrades to "most recent trace" on older MLflow versions that do not tag the working + directory. + """ candidates = [t for t in traces if trace_name(t) not in companion_names] + if working_directory: + scoped = [ + t for t in candidates + if _same_directory(trace_working_directory(t), working_directory) + ] + candidates = scoped or candidates + conversation = [t for t in candidates if trace_name(t) == SESSION_TRACE_NAME] + return conversation or candidates + + +def select_session_trace(traces: list, working_directory: str | None = None): + """Return the most recent trace for the current session, or None.""" + candidates = session_candidates(traces, working_directory) if not candidates: return None return max(candidates, key=_request_time_ms) @@ -134,10 +183,12 @@ def _resolve_trace_id(args) -> str: import mlflow - traces = mlflow.search_traces(max_results=50, return_type="list") - trace = select_most_recent(traces) + traces = mlflow.search_traces(max_results=SEARCH_WINDOW, return_type="list") + trace = select_session_trace(traces, working_directory=str(Path.cwd())) if trace is None: - raise RuntimeError("No traces found to annotate. Run a session first, or pass --trace-id.") + raise RuntimeError( + "No traces found to annotate. Run a coding session first, or pass --trace-id." + ) return trace.info.trace_id @@ -179,13 +230,16 @@ def cmd_list(args) -> None: import mlflow _configure_mlflow() - traces = mlflow.search_traces(max_results=args.limit, return_type="list") - traces = [t for t in traces if trace_name(t) not in COMPANION_TRACE_NAMES] - if not traces: + # Fetch more than the display limit so project scoping still yields up to --limit rows. + window = max(args.limit, SEARCH_WINDOW) + traces = mlflow.search_traces(max_results=window, return_type="list") + candidates = session_candidates(traces, working_directory=str(Path.cwd())) + candidates.sort(key=_request_time_ms, reverse=True) + candidates = candidates[: args.limit] + if not candidates: print("No traces found.") return - traces.sort(key=_request_time_ms, reverse=True) - print(_format_table(traces)) + print(_format_table(candidates)) def main() -> None: