feat(annotate-mlflow-trace): add trace tagging and feedback skill

.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"

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
 

annotate-mlflow-trace/SKILL.md

@@ -0,0 +1,82 @@
+---
+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 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
+
+```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 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: 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: current project's 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.
+
+## 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`.

annotate-mlflow-trace/scripts/test_trace_annotate.py

@@ -0,0 +1,133 @@
+#!/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, 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,
+        )
+    )
+
+
+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_session_trace_picks_latest():
+    traces = [_trace("tr-old", request_time_ms=100), _trace("tr-new", request_time_ms=200)]
+    assert trace_annotate.select_session_trace(traces).info.trace_id == "tr-new"
+
+
+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_session_trace(traces).info.trace_id == "tr-session"
+
+
+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():
+    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())