chore: import upstream snapshot with attribution

tests/claude_code/test_autolog.py

@@ -0,0 +1,137 @@
+import sys
+from unittest.mock import MagicMock, patch
+
+import pytest
+from claude_agent_sdk.types import AssistantMessage, ResultMessage, TextBlock, UserMessage
+
+import mlflow.anthropic
+from mlflow.anthropic.autolog import patched_claude_sdk_init
+
+
+def test_anthropic_autolog_without_claude_sdk():
+    sys.modules.pop("claude_agent_sdk", None)
+
+    with (
+        patch.dict(
+            "sys.modules",
+            {
+                "anthropic": MagicMock(__version__="0.35.0"),
+                "anthropic.resources": MagicMock(Messages=MagicMock, AsyncMessages=MagicMock),
+            },
+        ),
+        patch("mlflow.anthropic.safe_patch"),
+    ):
+        mlflow.anthropic.autolog()
+
+
+def _patch_sdk_init(mock_self, response_messages):
+    original_init = MagicMock()
+
+    async def fake_receive_response():
+        for msg in response_messages:
+            yield msg
+
+    mock_self.receive_response = fake_receive_response
+    patched_claude_sdk_init(original_init, mock_self)
+    return original_init
+
+
+def test_patched_claude_sdk_init_wraps_receive_response():
+    mock_self = MagicMock()
+
+    async def fake_receive_response():
+        yield "msg1"
+
+    mock_self.receive_response = fake_receive_response
+    original_init = MagicMock()
+    patched_claude_sdk_init(original_init, mock_self)
+
+    original_init.assert_called_once_with(mock_self, None)
+    assert mock_self.receive_response is not fake_receive_response
+
+
+@pytest.mark.asyncio
+async def test_receive_response_builds_trace():
+    mock_self = MagicMock()
+    messages = [
+        UserMessage(content="Hello"),
+        AssistantMessage(content=[TextBlock(text="Hi!")], model="claude-sonnet-4-20250514"),
+        ResultMessage(
+            subtype="success",
+            duration_ms=5000,
+            duration_api_ms=4000,
+            is_error=False,
+            num_turns=1,
+            session_id="test-session",
+            usage={"input_tokens": 100, "output_tokens": 20},
+        ),
+    ]
+    _patch_sdk_init(mock_self, messages)
+
+    with (
+        patch("mlflow.utils.autologging_utils.autologging_is_disabled", return_value=False),
+        patch("mlflow.claude_code.tracing.process_sdk_messages") as mock_process,
+    ):
+        [msg async for msg in mock_self.receive_response()]
+
+    mock_process.assert_called_once()
+    called_messages = mock_process.call_args[0][0]
+    assert len(called_messages) == 3
+    result_messages = [m for m in called_messages if isinstance(m, ResultMessage)]
+    assert len(result_messages) == 1
+    assert result_messages[0].usage == {"input_tokens": 100, "output_tokens": 20}
+
+
+@pytest.mark.asyncio
+async def test_query_captures_async_generator_prompt():
+    mock_self = MagicMock()
+
+    async def fake_query(prompt, *args, **kwargs):
+        # Consume the generator like the real SDK would
+        async for _ in prompt:
+            pass
+
+    mock_self.query = fake_query
+
+    response_messages = [
+        AssistantMessage(content=[TextBlock(text="Hi!")], model="claude-sonnet-4-20250514"),
+        ResultMessage(
+            subtype="success",
+            duration_ms=1000,
+            duration_api_ms=800,
+            is_error=False,
+            num_turns=1,
+            session_id="s",
+        ),
+    ]
+    _patch_sdk_init(mock_self, response_messages)
+
+    async def prompt_generator():
+        yield {"type": "user", "message": {"role": "user", "content": "Hello from generator"}}
+
+    with (
+        patch("mlflow.utils.autologging_utils.autologging_is_disabled", return_value=False),
+        patch("mlflow.claude_code.tracing.process_sdk_messages") as mock_process,
+    ):
+        await mock_self.query(prompt_generator())
+        [msg async for msg in mock_self.receive_response()]
+
+    mock_process.assert_called_once()
+    called_messages = mock_process.call_args[0][0]
+    user_messages = [m for m in called_messages if isinstance(m, UserMessage)]
+    assert len(user_messages) == 1
+    assert user_messages[0].content == "Hello from generator"
+
+
+@pytest.mark.asyncio
+async def test_receive_response_skips_when_autologging_disabled():
+    mock_self = MagicMock()
+    _patch_sdk_init(mock_self, ["msg1", "msg2"])
+
+    with (
+        patch("mlflow.utils.autologging_utils.autologging_is_disabled", return_value=True),
+        patch("mlflow.claude_code.tracing.process_sdk_messages") as mock_process,
+    ):
+        [msg async for msg in mock_self.receive_response()]
+
+    mock_process.assert_not_called()

tests/claude_code/test_cli.py

@@ -0,0 +1,224 @@
+import json
+from pathlib import Path
+from unittest import mock
+
+import pytest
+from click.testing import CliRunner
+
+from mlflow.claude_code.cli import commands
+
+
+@pytest.fixture
+def runner():
+    return CliRunner()
+
+
+@pytest.fixture(autouse=True)
+def _clear_mlflow_env(monkeypatch):
+    for name in (
+        "MLFLOW_TRACKING_URI",
+        "MLFLOW_EXPERIMENT_ID",
+        "MLFLOW_EXPERIMENT_NAME",
+    ):
+        monkeypatch.delenv(name, raising=False)
+
+
+def test_claude_help_command(runner):
+    result = runner.invoke(commands, ["--help"])
+    assert result.exit_code == 0
+    assert "Commands for autologging with MLflow" in result.output
+    assert "claude" in result.output
+
+
+def test_trace_command_help(runner):
+    result = runner.invoke(commands, ["claude", "--help"])
+    assert result.exit_code == 0
+    assert "Set up Claude Code tracing" in result.output
+    assert "--tracking-uri" in result.output
+    assert "--experiment-id" in result.output
+    assert "--non-interactive" in result.output
+    assert "--disable" in result.output
+    assert "--status" in result.output
+
+
+def test_trace_status_with_no_config(runner):
+    with runner.isolated_filesystem():
+        result = runner.invoke(commands, ["claude", "--status"])
+        assert result.exit_code == 0
+        assert "Claude tracing is not enabled" in result.output
+
+
+def test_trace_disable_with_no_config(runner):
+    with runner.isolated_filesystem():
+        result = runner.invoke(commands, ["claude", "--disable"])
+        assert result.exit_code == 0
+
+
+def test_claude_setup_installs_plugin_and_writes_env(runner):
+    with (
+        runner.isolated_filesystem(),
+        mock.patch("mlflow.claude_code.cli.ensure_plugin_installed") as mock_install,
+    ):
+        result = runner.invoke(commands, ["claude", "-u", "http://localhost:5000", "-e", "123"])
+        assert result.exit_code == 0
+
+        mock_install.assert_called_once_with(Path(".").resolve())
+
+        config = json.loads(Path(".claude/settings.json").read_text())
+        assert config["env"]["MLFLOW_CLAUDE_TRACING_ENABLED"] == "true"
+        assert config["env"]["MLFLOW_TRACKING_URI"] == "http://localhost:5000"
+        assert config["env"]["MLFLOW_EXPERIMENT_ID"] == "123"
+        assert "hooks" not in config
+
+
+def test_claude_setup_prompts_for_missing_values_in_interactive_mode(runner):
+    with (
+        runner.isolated_filesystem(),
+        mock.patch("mlflow.claude_code.cli.ensure_plugin_installed"),
+        mock.patch("mlflow.claude_code.cli._is_interactive_shell", return_value=True),
+        mock.patch("mlflow.get_tracking_uri", return_value="http://localhost:5000"),
+    ):
+        result = runner.invoke(commands, ["claude"], input="\n42\n")
+        assert result.exit_code == 0
+
+        config = json.loads(Path(".claude/settings.json").read_text())
+        assert config["env"]["MLFLOW_TRACKING_URI"] == "http://localhost:5000"
+        assert config["env"]["MLFLOW_EXPERIMENT_ID"] == "42"
+        assert "interactive mode" in result.output
+        assert "MLFLOW_TRACKING_URI and MLFLOW_EXPERIMENT_ID" in result.output
+
+
+def test_claude_setup_shows_plugin_install_message(runner):
+    with (
+        runner.isolated_filesystem(),
+        mock.patch("mlflow.claude_code.cli.ensure_plugin_installed"),
+    ):
+        result = runner.invoke(commands, ["claude", "-u", "http://localhost:5000", "-e", "123"])
+        assert result.exit_code == 0
+        assert "MLflow Claude plugin for Claude Code" in result.output
+        assert "Claude Code plugin installed" in result.output
+
+
+def test_claude_setup_non_interactive_uses_defaults(runner):
+    with (
+        runner.isolated_filesystem(),
+        mock.patch("mlflow.claude_code.cli.ensure_plugin_installed"),
+        mock.patch("mlflow.get_tracking_uri", return_value="file:///tmp/mlruns"),
+    ):
+        result = runner.invoke(commands, ["claude", "--non-interactive"])
+        assert result.exit_code == 0
+
+        config = json.loads(Path(".claude/settings.json").read_text())
+        assert config["env"]["MLFLOW_TRACKING_URI"] == "file:///tmp/mlruns"
+        assert config["env"]["MLFLOW_EXPERIMENT_ID"] == "0"
+
+
+def test_mlflow_cmd_empty_string_raises_error(runner):
+    with runner.isolated_filesystem():
+        result = runner.invoke(commands, ["claude", "--mlflow-cmd", ""])
+        assert result.exit_code != 0
+        assert "must not be empty or whitespace-only" in result.output
+
+
+def test_claude_setup_surfaces_plugin_install_failure(runner):
+    with (
+        runner.isolated_filesystem(),
+        mock.patch(
+            "mlflow.claude_code.cli.ensure_plugin_installed",
+            side_effect=RuntimeError("boom"),
+        ),
+    ):
+        result = runner.invoke(commands, ["claude"])
+        assert result.exit_code != 0
+        assert "boom" in result.output
+
+
+def test_mlflow_cmd_whitespace_only_raises_error(runner):
+    with runner.isolated_filesystem():
+        result = runner.invoke(commands, ["claude", "--mlflow-cmd", "   "])
+        assert result.exit_code != 0
+        assert "must not be empty or whitespace-only" in result.output
+
+
+def test_setup_rejects_experiment_id_and_name_together(runner):
+    with runner.isolated_filesystem():
+        result = runner.invoke(
+            commands,
+            ["claude", "--experiment-id", "1", "--experiment-name", "my-exp"],
+        )
+        assert result.exit_code != 0
+        assert "Choose either --experiment-id or --experiment-name" in result.output
+
+
+def test_claude_setup_with_local_flag(runner, monkeypatch):
+    monkeypatch.delenv("UV", raising=False)
+    monkeypatch.delenv("PIXI_ENVIRONMENT_NAME", raising=False)
+
+    with (
+        runner.isolated_filesystem(),
+        mock.patch("mlflow.claude_code.cli.ensure_plugin_installed"),
+    ):
+        result = runner.invoke(commands, ["claude", "--local"])
+        assert result.exit_code == 0
+
+        local_path = Path(".claude/settings.local.json")
+        assert local_path.exists()
+
+        with open(local_path) as f:
+            config = json.load(f)
+        assert "MLFLOW_CLAUDE_TRACING_ENABLED" in config.get("env", {})
+
+        settings_path = Path(".claude/settings.json")
+        assert not settings_path.exists()
+
+
+def test_claude_setup_local_status(runner, monkeypatch):
+    monkeypatch.delenv("UV", raising=False)
+    monkeypatch.delenv("PIXI_ENVIRONMENT_NAME", raising=False)
+
+    with (
+        runner.isolated_filesystem(),
+        mock.patch("mlflow.claude_code.cli.ensure_plugin_installed"),
+    ):
+        result = runner.invoke(commands, ["claude", "--local"])
+        assert result.exit_code == 0
+
+        result = runner.invoke(commands, ["claude", "--status"])
+        assert result.exit_code == 0
+        assert "Claude tracing is enabled" in result.output
+
+
+def test_claude_disable_cleans_local_without_flag(runner, monkeypatch):
+    monkeypatch.delenv("UV", raising=False)
+    monkeypatch.delenv("PIXI_ENVIRONMENT_NAME", raising=False)
+
+    with (
+        runner.isolated_filesystem(),
+        mock.patch("mlflow.claude_code.cli.ensure_plugin_installed"),
+    ):
+        result = runner.invoke(commands, ["claude", "--local"])
+        assert result.exit_code == 0
+
+        # Disable without --local should still clean settings.local.json
+        result = runner.invoke(commands, ["claude", "--disable"])
+        assert result.exit_code == 0
+        assert "Claude tracing disabled" in result.output
+
+
+def test_local_with_status_raises_error(runner):
+    result = runner.invoke(commands, ["claude", "--local", "--status"])
+    assert result.exit_code != 0
+    assert "--local can only be used during setup" in result.output
+
+
+def test_local_with_disable_raises_error(runner):
+    result = runner.invoke(commands, ["claude", "--local", "--disable"])
+    assert result.exit_code != 0
+    assert "--local can only be used during setup" in result.output
+
+
+def test_stop_hook_subcommand_is_routable(runner):
+    with mock.patch("mlflow.claude_code.cli.stop_hook_handler") as mock_handler:
+        result = runner.invoke(commands, ["claude", "stop-hook"])
+        assert result.exit_code == 0
+        mock_handler.assert_called_once()

tests/claude_code/test_config.py

@@ -0,0 +1,333 @@
+import json
+
+import pytest
+
+from mlflow.claude_code.config import (
+    MLFLOW_TRACING_ENABLED,
+    get_env_var,
+    get_tracing_status,
+    load_claude_config,
+    save_claude_config,
+    setup_environment_config,
+)
+
+
+@pytest.fixture
+def temp_settings_path(tmp_path):
+    """Provide a temporary settings.json path for tests."""
+    return tmp_path / "settings.json"
+
+
+def test_load_claude_config_valid_json(temp_settings_path):
+    config_data = {"tools": {"computer_20241022": {"name": "computer"}}}
+    with open(temp_settings_path, "w") as f:
+        json.dump(config_data, f)
+
+    result = load_claude_config(temp_settings_path)
+    assert result == config_data
+
+
+def test_load_claude_config_missing_file(tmp_path):
+    non_existent_path = tmp_path / "non_existent.json"
+    result = load_claude_config(non_existent_path)
+    assert result == {}
+
+
+def test_load_claude_config_invalid_json(temp_settings_path):
+    with open(temp_settings_path, "w") as f:
+        f.write("invalid json content")
+
+    result = load_claude_config(temp_settings_path)
+    assert result == {}
+
+
+def test_save_claude_config_creates_file(temp_settings_path):
+    config_data = {"test": "value"}
+    save_claude_config(temp_settings_path, config_data)
+
+    assert temp_settings_path.exists()
+    saved_data = json.loads(temp_settings_path.read_text())
+    assert saved_data == config_data
+
+
+def test_save_claude_config_creates_directory(tmp_path):
+    nested_path = tmp_path / "nested" / "dir" / "settings.json"
+    config_data = {"test": "value"}
+
+    save_claude_config(nested_path, config_data)
+
+    assert nested_path.exists()
+    saved_data = json.loads(nested_path.read_text())
+    assert saved_data == config_data
+
+
+def test_get_env_var_from_os_environment_when_no_settings(tmp_path, monkeypatch):
+    monkeypatch.setenv(MLFLOW_TRACING_ENABLED, "test_os_value")
+    monkeypatch.chdir(tmp_path)
+
+    result = get_env_var(MLFLOW_TRACING_ENABLED, "default")
+    assert result == "test_os_value"
+
+
+def test_get_env_var_os_env_takes_precedence_over_settings(tmp_path, monkeypatch):
+    monkeypatch.setenv(MLFLOW_TRACING_ENABLED, "os_value")
+
+    config_data = {"env": {MLFLOW_TRACING_ENABLED: "settings_value"}}
+    claude_settings_path = tmp_path / ".claude" / "settings.json"
+    claude_settings_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(claude_settings_path, "w") as f:
+        json.dump(config_data, f)
+
+    monkeypatch.chdir(tmp_path)
+    result = get_env_var(MLFLOW_TRACING_ENABLED, "default")
+    assert result == "os_value"
+
+
+def test_get_env_var_falls_back_to_os_env_when_not_in_settings(tmp_path, monkeypatch):
+    monkeypatch.setenv(MLFLOW_TRACING_ENABLED, "os_value")
+
+    config_data = {"env": {"OTHER_VAR": "other_value"}}
+    claude_settings_path = tmp_path / ".claude" / "settings.json"
+    claude_settings_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(claude_settings_path, "w") as f:
+        json.dump(config_data, f)
+
+    monkeypatch.chdir(tmp_path)
+    result = get_env_var(MLFLOW_TRACING_ENABLED, "default")
+    assert result == "os_value"
+
+
+def test_get_env_var_from_settings_local_json(tmp_path, monkeypatch):
+    monkeypatch.delenv(MLFLOW_TRACING_ENABLED, raising=False)
+
+    config_data = {"env": {MLFLOW_TRACING_ENABLED: "local_value"}}
+    claude_dir = tmp_path / ".claude"
+    claude_dir.mkdir(parents=True, exist_ok=True)
+    with open(claude_dir / "settings.local.json", "w") as f:
+        json.dump(config_data, f)
+
+    monkeypatch.chdir(tmp_path)
+    result = get_env_var(MLFLOW_TRACING_ENABLED, "default")
+    assert result == "local_value"
+
+
+def test_get_env_var_settings_local_overrides_settings(tmp_path, monkeypatch):
+    monkeypatch.delenv(MLFLOW_TRACING_ENABLED, raising=False)
+
+    claude_dir = tmp_path / ".claude"
+    claude_dir.mkdir(parents=True, exist_ok=True)
+
+    with open(claude_dir / "settings.json", "w") as f:
+        json.dump({"env": {MLFLOW_TRACING_ENABLED: "shared_value"}}, f)
+    with open(claude_dir / "settings.local.json", "w") as f:
+        json.dump({"env": {MLFLOW_TRACING_ENABLED: "local_value"}}, f)
+
+    monkeypatch.chdir(tmp_path)
+    result = get_env_var(MLFLOW_TRACING_ENABLED, "default")
+    assert result == "local_value"
+
+
+def test_get_env_var_falls_through_local_to_settings(tmp_path, monkeypatch):
+    monkeypatch.delenv("MLFLOW_TRACKING_URI", raising=False)
+
+    claude_dir = tmp_path / ".claude"
+    claude_dir.mkdir(parents=True, exist_ok=True)
+
+    # local has ENABLED, shared has URI
+    with open(claude_dir / "settings.local.json", "w") as f:
+        json.dump({"env": {MLFLOW_TRACING_ENABLED: "true"}}, f)
+    with open(claude_dir / "settings.json", "w") as f:
+        json.dump({"env": {"MLFLOW_TRACKING_URI": "https://example.com"}}, f)
+
+    monkeypatch.chdir(tmp_path)
+    assert get_env_var(MLFLOW_TRACING_ENABLED, "default") == "true"
+    assert get_env_var("MLFLOW_TRACKING_URI", "default") == "https://example.com"
+
+
+def test_get_env_var_os_env_takes_precedence_over_local(tmp_path, monkeypatch):
+    monkeypatch.setenv(MLFLOW_TRACING_ENABLED, "os_value")
+
+    claude_dir = tmp_path / ".claude"
+    claude_dir.mkdir(parents=True, exist_ok=True)
+    with open(claude_dir / "settings.local.json", "w") as f:
+        json.dump({"env": {MLFLOW_TRACING_ENABLED: "local_value"}}, f)
+
+    monkeypatch.chdir(tmp_path)
+    result = get_env_var(MLFLOW_TRACING_ENABLED, "default")
+    assert result == "os_value"
+
+
+def test_get_tracing_status_merges_local_and_shared(tmp_path):
+    claude_dir = tmp_path / ".claude"
+    claude_dir.mkdir(parents=True, exist_ok=True)
+    settings_path = claude_dir / "settings.json"
+
+    # URI in shared, ENABLED in local
+    with open(settings_path, "w") as f:
+        json.dump({"env": {"MLFLOW_TRACKING_URI": "https://example.com"}}, f)
+    with open(claude_dir / "settings.local.json", "w") as f:
+        json.dump({"env": {MLFLOW_TRACING_ENABLED: "true"}}, f)
+
+    status = get_tracing_status(settings_path)
+    assert status.enabled is True
+    assert status.tracking_uri == "https://example.com"
+
+
+def test_get_tracing_status_local_overrides_shared(tmp_path):
+    claude_dir = tmp_path / ".claude"
+    claude_dir.mkdir(parents=True, exist_ok=True)
+    settings_path = claude_dir / "settings.json"
+
+    with open(settings_path, "w") as f:
+        json.dump({"env": {MLFLOW_TRACING_ENABLED: "true"}}, f)
+    with open(claude_dir / "settings.local.json", "w") as f:
+        json.dump({"env": {MLFLOW_TRACING_ENABLED: "false"}}, f)
+
+    status = get_tracing_status(settings_path)
+    assert status.enabled is False
+
+
+def test_get_tracing_status_from_local_only(tmp_path):
+    claude_dir = tmp_path / ".claude"
+    claude_dir.mkdir(parents=True, exist_ok=True)
+    settings_path = claude_dir / "settings.json"
+
+    with open(claude_dir / "settings.local.json", "w") as f:
+        json.dump({"env": {MLFLOW_TRACING_ENABLED: "true"}}, f)
+
+    status = get_tracing_status(settings_path)
+    assert status.enabled is True
+
+
+def test_get_env_var_default_when_not_found(tmp_path, monkeypatch):
+    # Ensure OS env var is not set
+    monkeypatch.delenv(MLFLOW_TRACING_ENABLED, raising=False)
+
+    # Create empty settings file in .claude directory
+    claude_settings_path = tmp_path / ".claude" / "settings.json"
+    claude_settings_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(claude_settings_path, "w") as f:
+        json.dump({}, f)
+
+    # Change to temp directory so .claude/settings.json is found
+    monkeypatch.chdir(tmp_path)
+    result = get_env_var(MLFLOW_TRACING_ENABLED, "default_value")
+    assert result == "default_value"
+
+
+def test_get_tracing_status_enabled(temp_settings_path):
+    # Create settings with tracing enabled
+    config_data = {"env": {MLFLOW_TRACING_ENABLED: "true"}}
+    with open(temp_settings_path, "w") as f:
+        json.dump(config_data, f)
+
+    status = get_tracing_status(temp_settings_path)
+    assert status.enabled is True
+    assert hasattr(status, "tracking_uri")
+
+
+def test_get_tracing_status_disabled(temp_settings_path):
+    # Create settings with tracing disabled
+    config_data = {"env": {MLFLOW_TRACING_ENABLED: "false"}}
+    with open(temp_settings_path, "w") as f:
+        json.dump(config_data, f)
+
+    status = get_tracing_status(temp_settings_path)
+    assert status.enabled is False
+
+
+def test_get_tracing_status_no_config(tmp_path):
+    non_existent_path = tmp_path / "missing.json"
+    status = get_tracing_status(non_existent_path)
+    assert status.enabled is False
+    assert status.reason == "No configuration found"
+
+
+def test_setup_environment_config_new_file(temp_settings_path):
+    tracking_uri = "test://localhost"
+    experiment_id = "123"
+
+    setup_environment_config(temp_settings_path, tracking_uri, experiment_id)
+
+    # Verify file was created
+    assert temp_settings_path.exists()
+
+    # Verify configuration contents
+    config = json.loads(temp_settings_path.read_text())
+
+    env_vars = config["env"]
+    assert env_vars[MLFLOW_TRACING_ENABLED] == "true"
+    assert env_vars["MLFLOW_TRACKING_URI"] == tracking_uri
+    assert env_vars["MLFLOW_EXPERIMENT_ID"] == experiment_id
+
+
+def test_setup_environment_config_experiment_id_precedence(temp_settings_path):
+    # Create existing config with different experiment ID
+    existing_config = {
+        "env": {
+            MLFLOW_TRACING_ENABLED: "true",
+            "MLFLOW_EXPERIMENT_ID": "old_id",
+            "MLFLOW_TRACKING_URI": "old_uri",
+        }
+    }
+    with open(temp_settings_path, "w") as f:
+        json.dump(existing_config, f)
+
+    new_tracking_uri = "new://localhost"
+    new_experiment_id = "new_id"
+
+    setup_environment_config(temp_settings_path, new_tracking_uri, new_experiment_id)
+
+    # Verify configuration was updated
+    config = json.loads(temp_settings_path.read_text())
+
+    env_vars = config["env"]
+    assert env_vars[MLFLOW_TRACING_ENABLED] == "true"
+    assert env_vars["MLFLOW_TRACKING_URI"] == new_tracking_uri
+    assert env_vars["MLFLOW_EXPERIMENT_ID"] == new_experiment_id
+
+
+def test_setup_environment_config_defaults_to_current_tracking_uri_and_default_experiment(
+    temp_settings_path, monkeypatch
+):
+    monkeypatch.delenv("MLFLOW_TRACKING_URI", raising=False)
+    monkeypatch.delenv("MLFLOW_EXPERIMENT_ID", raising=False)
+    monkeypatch.delenv("MLFLOW_EXPERIMENT_NAME", raising=False)
+    monkeypatch.setattr("mlflow.get_tracking_uri", lambda: "file:///tmp/mlruns")
+
+    setup_environment_config(temp_settings_path)
+
+    config = json.loads(temp_settings_path.read_text())
+    env_vars = config["env"]
+    assert env_vars["MLFLOW_TRACKING_URI"] == "file:///tmp/mlruns"
+    assert env_vars["MLFLOW_EXPERIMENT_ID"] == "0"
+
+
+def test_setup_environment_config_resolves_experiment_name_to_id(temp_settings_path, monkeypatch):
+    class DummyExperiment:
+        experiment_id = "456"
+
+    class DummyClient:
+        def __init__(self, tracking_uri):
+            assert tracking_uri == "http://localhost:5000"
+
+        def get_experiment_by_name(self, name):
+            assert name == "my-exp"
+            return DummyExperiment()
+
+        def create_experiment(self, name):
+            raise AssertionError("create_experiment should not be called when experiment exists")
+
+    monkeypatch.setattr("mlflow.tracking.client.MlflowClient", DummyClient)
+
+    setup_environment_config(
+        temp_settings_path,
+        tracking_uri="http://localhost:5000",
+        experiment_name="my-exp",
+    )
+
+    config = json.loads(temp_settings_path.read_text())
+    env_vars = config["env"]
+    assert env_vars["MLFLOW_TRACKING_URI"] == "http://localhost:5000"
+    assert env_vars["MLFLOW_EXPERIMENT_ID"] == "456"
+    assert env_vars["MLFLOW_EXPERIMENT_NAME"] == "my-exp"

tests/claude_code/test_plugin.py

@@ -0,0 +1,49 @@
+import json
+import subprocess
+from unittest import mock
+
+import click
+import pytest
+
+from mlflow.claude_code.plugin import disable_tracing_plugin, ensure_plugin_installed
+
+
+def test_disable_tracing_plugin_removes_env_only(tmp_path):
+    settings_path = tmp_path / ".claude" / "settings.json"
+    settings_path.parent.mkdir(parents=True)
+    settings_path.write_text(
+        json.dumps({
+            "env": {
+                "MLFLOW_CLAUDE_TRACING_ENABLED": "true",
+                "MLFLOW_TRACKING_URI": "http://localhost:5000",
+                "MLFLOW_EXPERIMENT_ID": "123",
+            },
+            "other": "keep-me",
+        })
+    )
+
+    assert disable_tracing_plugin(settings_path) is True
+    config = json.loads(settings_path.read_text())
+    assert config == {"other": "keep-me"}
+
+
+def test_ensure_plugin_installed_runs_marketplace_add_and_install(tmp_path):
+    completed = subprocess.CompletedProcess(args=[], returncode=0, stdout="", stderr="")
+
+    with (
+        mock.patch("shutil.which", return_value="/usr/local/bin/claude"),
+        mock.patch("subprocess.run", return_value=completed) as mock_run,
+    ):
+        ensure_plugin_installed(tmp_path)
+
+    assert mock_run.call_count == 2
+    first_command = mock_run.call_args_list[0].args[0]
+    second_command = mock_run.call_args_list[1].args[0]
+    assert first_command[:5] == ["claude", "plugin", "marketplace", "add", "mlflow/mlflow"]
+    assert second_command[:4] == ["claude", "plugin", "install", "mlflow-tracing@mlflow-plugins"]
+
+
+def test_ensure_plugin_installed_requires_claude_binary(tmp_path):
+    with mock.patch("shutil.which", return_value=None):
+        with pytest.raises(click.ClickException, match="Claude Code CLI"):
+            ensure_plugin_installed(tmp_path)

tests/claude_code/test_tracing.py

@@ -0,0 +1,905 @@
+import importlib
+import json
+import logging
+from pathlib import Path
+
+import pytest
+from claude_agent_sdk.types import (
+    AssistantMessage,
+    ResultMessage,
+    TextBlock,
+    ToolResultBlock,
+    ToolUseBlock,
+    UserMessage,
+)
+
+import mlflow
+import mlflow.claude_code.tracing as tracing_module
+from mlflow.claude_code.tracing import (
+    CLAUDE_TRACING_LEVEL,
+    METADATA_KEY_CLAUDE_CODE_VERSION,
+    find_last_user_message_index,
+    get_hook_response,
+    parse_timestamp_to_ns,
+    process_sdk_messages,
+    process_transcript,
+    setup_logging,
+)
+from mlflow.entities.span import SpanType
+from mlflow.tracing.constant import SpanAttributeKey, TraceMetadataKey
+
+# ============================================================================
+# TIMESTAMP PARSING TESTS
+# ============================================================================
+
+
+def test_parse_timestamp_to_ns_iso_string():
+    iso_timestamp = "2024-01-15T10:30:45.123456Z"
+    result = parse_timestamp_to_ns(iso_timestamp)
+
+    # Verify it returns an integer (nanoseconds)
+    assert isinstance(result, int)
+    assert result > 0
+
+
+def test_parse_timestamp_to_ns_unix_seconds():
+    unix_timestamp = 1705312245.123456
+    result = parse_timestamp_to_ns(unix_timestamp)
+
+    # Should convert seconds to nanoseconds
+    expected = int(unix_timestamp * 1_000_000_000)
+    assert result == expected
+
+
+def test_parse_timestamp_to_ns_large_number():
+    large_timestamp = 1705312245123
+    result = parse_timestamp_to_ns(large_timestamp)
+
+    # Function treats large numbers as seconds and converts to nanoseconds
+    # Just verify we get a reasonable nanosecond value
+    assert isinstance(result, int)
+    assert result > 0
+
+
+# ============================================================================
+# LOGGING TESTS
+# ============================================================================
+
+
+def test_setup_logging_creates_logger(monkeypatch, tmp_path):
+    monkeypatch.chdir(tmp_path)
+    logger = setup_logging()
+
+    # Verify logger was created
+    assert logger is not None
+    assert logger.name == "mlflow.claude_code.tracing"
+
+    # Verify log directory was created
+    log_dir = tmp_path / ".claude" / "mlflow"
+    assert log_dir.exists()
+    assert log_dir.is_dir()
+
+
+def test_custom_logging_level():
+    setup_logging()
+
+    assert CLAUDE_TRACING_LEVEL > logging.INFO
+    assert CLAUDE_TRACING_LEVEL < logging.WARNING
+    assert logging.getLevelName(CLAUDE_TRACING_LEVEL) == "CLAUDE_TRACING"
+
+
+def test_get_logger_lazy_initialization(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+    monkeypatch.chdir(tmp_path)
+
+    # Force reload to reset the module state
+    importlib.reload(tracing_module)
+
+    log_dir = tmp_path / ".claude" / "mlflow"
+
+    # Before calling get_logger(), the log directory should NOT exist
+    assert not log_dir.exists()
+
+    # Call get_logger() for the first time - this should trigger initialization
+    logger1 = tracing_module.get_logger()
+
+    # After calling get_logger(), the log directory SHOULD exist
+    assert log_dir.exists()
+    assert log_dir.is_dir()
+
+    # Verify logger was created properly
+    assert logger1 is not None
+    assert logger1.name == "mlflow.claude_code.tracing"
+
+    # Call get_logger() again - should return the same logger instance
+    logger2 = tracing_module.get_logger()
+    assert logger2 is logger1
+
+
+# ============================================================================
+# HOOK RESPONSE TESTS
+# ============================================================================
+
+
+def test_get_hook_response_success():
+    response = get_hook_response()
+    assert response == {"continue": True}
+
+
+def test_get_hook_response_with_error():
+    response = get_hook_response(error="Test error")
+    assert response == {"continue": False, "stopReason": "Test error"}
+
+
+def test_get_hook_response_with_additional_fields():
+    response = get_hook_response(custom_field="value")
+    assert response == {"continue": True, "custom_field": "value"}
+
+
+# ============================================================================
+# ASYNC TRACE LOGGING UTILITY TESTS
+# ============================================================================
+
+
+def test_flush_trace_async_logging_calls_flush(monkeypatch):
+    mock_exporter = type("MockExporter", (), {"_async_queue": True})()
+    monkeypatch.setattr(tracing_module, "_get_trace_exporter", lambda: mock_exporter)
+    flushed = []
+    monkeypatch.setattr(mlflow, "flush_trace_async_logging", lambda: flushed.append(True))
+    tracing_module._flush_trace_async_logging()
+    assert len(flushed) == 1
+
+
+def test_flush_trace_async_logging_skips_without_async_queue(monkeypatch):
+    mock_exporter = object()  # no _async_queue attribute
+    monkeypatch.setattr(tracing_module, "_get_trace_exporter", lambda: mock_exporter)
+    flushed = []
+    monkeypatch.setattr(mlflow, "flush_trace_async_logging", lambda: flushed.append(True))
+    tracing_module._flush_trace_async_logging()
+    assert len(flushed) == 0
+
+
+# ============================================================================
+# INTEGRATION TESTS
+# ============================================================================
+
+# Sample Claude Code transcript for testing
+DUMMY_TRANSCRIPT = [
+    {
+        "type": "user",
+        "message": {"role": "user", "content": "What is 2 + 2?"},
+        "timestamp": "2025-01-15T10:00:00.000Z",
+        "sessionId": "test-session-123",
+    },
+    {
+        "type": "assistant",
+        "message": {
+            "role": "assistant",
+            "content": [{"type": "text", "text": "Let me calculate that for you."}],
+        },
+        "timestamp": "2025-01-15T10:00:01.000Z",
+    },
+    {
+        "type": "assistant",
+        "message": {
+            "role": "assistant",
+            "content": [
+                {
+                    "type": "tool_use",
+                    "id": "tool_123",
+                    "name": "Bash",
+                    "input": {"command": "echo $((2 + 2))"},
+                }
+            ],
+        },
+        "timestamp": "2025-01-15T10:00:02.000Z",
+    },
+    {
+        "type": "user",
+        "message": {
+            "role": "user",
+            "content": [{"type": "tool_result", "tool_use_id": "tool_123", "content": "4"}],
+        },
+        "timestamp": "2025-01-15T10:00:03.000Z",
+    },
+    {
+        "type": "assistant",
+        "message": {
+            "role": "assistant",
+            "content": [{"type": "text", "text": "The answer is 4."}],
+        },
+        "timestamp": "2025-01-15T10:00:04.000Z",
+    },
+]
+
+
+@pytest.fixture
+def mock_transcript_file(tmp_path):
+    transcript_path = tmp_path / "transcript.jsonl"
+    with open(transcript_path, "w") as f:
+        for entry in DUMMY_TRANSCRIPT:
+            f.write(json.dumps(entry) + "\n")
+    return str(transcript_path)
+
+
+def test_process_transript_creates_trace(mock_transcript_file):
+    trace = process_transcript(mock_transcript_file, "test-session-123")
+
+    # Verify trace was created
+    assert trace is not None
+
+    # Verify trace has spans
+    spans = list(trace.search_spans())
+    assert len(spans) > 0
+
+    # Verify root span and metadata
+    root_span = trace.data.spans[0]
+    assert root_span.name == "claude_code_conversation"
+    assert root_span.span_type == SpanType.AGENT
+    assert trace.info.trace_metadata.get("mlflow.trace.session") == "test-session-123"
+
+
+def test_process_transcript_creates_spans(mock_transcript_file):
+    trace = process_transcript(mock_transcript_file, "test-session-123")
+
+    assert trace is not None
+
+    # Verify trace has spans
+    spans = list(trace.search_spans())
+    assert len(spans) > 0
+
+    # Find LLM and tool spans
+    llm_spans = [s for s in spans if s.span_type == SpanType.LLM]
+    tool_spans = [s for s in spans if s.span_type == SpanType.TOOL]
+
+    assert len(llm_spans) == 2
+    assert len(tool_spans) == 1
+
+    # Verify tool span has proper attributes
+    tool_span = tool_spans[0]
+    assert tool_span.name == "tool_Bash"
+
+    # Verify LLM spans have MESSAGE_FORMAT set to "anthropic" for Chat UI rendering
+    for llm_span in llm_spans:
+        assert llm_span.get_attribute(SpanAttributeKey.MESSAGE_FORMAT) == "anthropic"
+
+    # Verify LLM span outputs are in Anthropic response format
+    first_llm = llm_spans[0]
+    outputs = first_llm.outputs
+    assert outputs["type"] == "message"
+    assert outputs["role"] == "assistant"
+    assert isinstance(outputs["content"], list)
+
+    # Verify LLM span inputs contain messages in Anthropic format
+    inputs = first_llm.inputs
+    assert "messages" in inputs
+    messages = inputs["messages"]
+    assert any(m["role"] == "user" for m in messages)
+
+
+def test_process_transcript_returns_none_for_nonexistent_file():
+    result = process_transcript("/nonexistent/path/transcript.jsonl", "test-session-123")
+    assert result is None
+
+
+def test_process_transcript_links_trace_to_run(mock_transcript_file):
+    with mlflow.start_run() as run:
+        trace = process_transcript(mock_transcript_file, "test-session-123")
+
+        assert trace is not None
+        assert trace.info.trace_metadata.get(TraceMetadataKey.SOURCE_RUN) == run.info.run_id
+
+
+# Sample Claude Code transcript with token usage for testing
+DUMMY_TRANSCRIPT_WITH_USAGE = [
+    {
+        "type": "user",
+        "message": {"role": "user", "content": "Hello Claude!"},
+        "timestamp": "2025-01-15T10:00:00.000Z",
+        "sessionId": "test-session-usage",
+    },
+    {
+        "type": "assistant",
+        "message": {
+            "role": "assistant",
+            "content": [{"type": "text", "text": "Hello! How can I help you today?"}],
+            "model": "claude-sonnet-4-20250514",
+            "usage": {"input_tokens": 150, "output_tokens": 25},
+        },
+        "timestamp": "2025-01-15T10:00:01.000Z",
+    },
+]
+
+
+@pytest.fixture
+def mock_transcript_file_with_usage(tmp_path):
+    transcript_path = tmp_path / "transcript_with_usage.jsonl"
+    with open(transcript_path, "w") as f:
+        for entry in DUMMY_TRANSCRIPT_WITH_USAGE:
+            f.write(json.dumps(entry) + "\n")
+    return str(transcript_path)
+
+
+def test_process_transcript_tracks_token_usage(mock_transcript_file_with_usage):
+    trace = process_transcript(mock_transcript_file_with_usage, "test-session-usage")
+
+    assert trace is not None
+
+    # Find the LLM span
+    spans = list(trace.search_spans())
+    llm_spans = [s for s in spans if s.span_type == SpanType.LLM]
+
+    assert len(llm_spans) == 1
+    llm_span = llm_spans[0]
+
+    # Verify token usage is tracked using the standardized CHAT_USAGE attribute
+    token_usage = llm_span.get_attribute(SpanAttributeKey.CHAT_USAGE)
+    assert token_usage is not None
+    assert token_usage["input_tokens"] == 150
+    assert token_usage["output_tokens"] == 25
+    assert token_usage["total_tokens"] == 175
+
+    # Verify trace-level token usage aggregation works
+    assert trace.info.token_usage is not None
+    assert trace.info.token_usage["input_tokens"] == 150
+    assert trace.info.token_usage["output_tokens"] == 25
+    assert trace.info.token_usage["total_tokens"] == 175
+
+
+def test_process_transcript_preserves_cache_tokens(tmp_path):
+    """Verify cache_read/cache_creation fields from Anthropic usage survive on the
+    CHAT_USAGE span attribute so prompt-cache hit rate is observable.
+    """
+    transcript_entries = [
+        {
+            "type": "user",
+            "message": {"role": "user", "content": "Cached prompt"},
+            "timestamp": "2025-01-15T10:00:00.000Z",
+            "sessionId": "cache-transcript-session",
+        },
+        {
+            "type": "assistant",
+            "message": {
+                "role": "assistant",
+                "content": [{"type": "text", "text": "Answer using cache."}],
+                "model": "claude-sonnet-4-20250514",
+                "usage": {
+                    "input_tokens": 36,
+                    "cache_creation_input_tokens": 23554,
+                    "cache_read_input_tokens": 139035,
+                    "output_tokens": 3344,
+                },
+            },
+            "timestamp": "2025-01-15T10:00:01.000Z",
+        },
+    ]
+
+    transcript_path = tmp_path / "transcript_cache.jsonl"
+    with open(transcript_path, "w") as f:
+        for entry in transcript_entries:
+            f.write(json.dumps(entry) + "\n")
+
+    trace = process_transcript(str(transcript_path), "cache-transcript-session")
+
+    assert trace is not None
+    llm_spans = [s for s in trace.search_spans() if s.span_type == SpanType.LLM]
+    assert len(llm_spans) == 1
+
+    # input_tokens is the non-cached input the Anthropic API reports, matching
+    # mlflow.anthropic.autolog. Cache fields are exposed as separate keys so
+    # consumers can compute cache hit rate.
+    token_usage = llm_spans[0].get_attribute(SpanAttributeKey.CHAT_USAGE)
+    assert token_usage["input_tokens"] == 36
+    assert token_usage["output_tokens"] == 3344
+    assert token_usage["total_tokens"] == 36 + 3344
+    assert token_usage["cache_read_input_tokens"] == 139035
+    assert token_usage["cache_creation_input_tokens"] == 23554
+
+
+# ============================================================================
+# SDK MESSAGE PROCESSING TESTS
+# ============================================================================
+
+
+def test_process_sdk_messages_empty_list():
+    assert process_sdk_messages([]) is None
+
+
+def test_process_sdk_messages_no_user_prompt():
+    messages = [
+        AssistantMessage(
+            content=[TextBlock(text="Hello!")],
+            model="claude-sonnet-4-20250514",
+        ),
+    ]
+    assert process_sdk_messages(messages) is None
+
+
+def test_process_sdk_messages_simple_conversation():
+    messages = [
+        UserMessage(content="What is 2 + 2?"),
+        AssistantMessage(
+            content=[TextBlock(text="The answer is 4.")],
+            model="claude-sonnet-4-20250514",
+        ),
+        ResultMessage(
+            subtype="success",
+            duration_ms=1000,
+            duration_api_ms=800,
+            is_error=False,
+            num_turns=1,
+            session_id="test-sdk-session",
+            usage={"input_tokens": 100, "output_tokens": 20},
+        ),
+    ]
+
+    trace = process_sdk_messages(messages, "test-sdk-session")
+
+    assert trace is not None
+    spans = list(trace.search_spans())
+
+    root_span = trace.data.spans[0]
+    assert root_span.name == "claude_code_conversation"
+    assert root_span.span_type == SpanType.AGENT
+
+    # LLM span should have conversation context as input in Anthropic format
+    llm_spans = [s for s in spans if s.span_type == SpanType.LLM]
+    assert len(llm_spans) == 1
+    assert llm_spans[0].name == "llm"
+    assert llm_spans[0].inputs["model"] == "claude-sonnet-4-20250514"
+    assert llm_spans[0].inputs["messages"] == [{"role": "user", "content": "What is 2 + 2?"}]
+    assert llm_spans[0].get_attribute(SpanAttributeKey.MESSAGE_FORMAT) == "anthropic"
+
+    # Output should be in Anthropic response format
+    outputs = llm_spans[0].outputs
+    assert outputs["type"] == "message"
+    assert outputs["role"] == "assistant"
+    assert outputs["content"] == [{"type": "text", "text": "The answer is 4."}]
+
+    # Token usage from ResultMessage should be on the root span and trace level
+    token_usage = root_span.get_attribute(SpanAttributeKey.CHAT_USAGE)
+    assert token_usage is not None
+    assert token_usage["input_tokens"] == 100
+    assert token_usage["output_tokens"] == 20
+    assert token_usage["total_tokens"] == 120
+
+    assert trace.info.token_usage is not None
+    assert trace.info.token_usage["input_tokens"] == 100
+    assert trace.info.token_usage["output_tokens"] == 20
+    assert trace.info.token_usage["total_tokens"] == 120
+
+    # Duration should reflect ResultMessage.duration_ms (1000ms = 1s)
+    duration_ns = root_span.end_time_ns - root_span.start_time_ns
+    assert abs(duration_ns - 1_000_000_000) < 1_000_000  # within 1ms tolerance
+
+    assert trace.info.trace_metadata.get("mlflow.trace.session") == "test-sdk-session"
+    assert trace.info.request_preview == "What is 2 + 2?"
+    assert trace.info.response_preview == "The answer is 4."
+
+
+def test_process_sdk_messages_multiple_tools():
+    messages = [
+        UserMessage(content="Read two files"),
+        AssistantMessage(
+            content=[
+                ToolUseBlock(id="tool_1", name="Read", input={"path": "a.py"}),
+                ToolUseBlock(id="tool_2", name="Read", input={"path": "b.py"}),
+            ],
+            model="claude-sonnet-4-20250514",
+        ),
+        UserMessage(
+            content=[
+                ToolResultBlock(tool_use_id="tool_1", content="content of a"),
+                ToolResultBlock(tool_use_id="tool_2", content="content of b"),
+            ],
+            tool_use_result={"tool_use_id": "tool_1"},
+        ),
+        AssistantMessage(
+            content=[TextBlock(text="Here are the contents.")],
+            model="claude-sonnet-4-20250514",
+        ),
+        ResultMessage(
+            subtype="success",
+            duration_ms=2000,
+            duration_api_ms=1500,
+            is_error=False,
+            num_turns=2,
+            session_id="multi-tool-session",
+        ),
+    ]
+
+    trace = process_sdk_messages(messages, "multi-tool-session")
+
+    assert trace is not None
+    spans = list(trace.search_spans())
+
+    tool_spans = [s for s in spans if s.span_type == SpanType.TOOL]
+    assert len(tool_spans) == 2
+    assert all(s.name == "tool_Read" for s in tool_spans)
+    tool_results = {s.outputs["result"] for s in tool_spans}
+    assert tool_results == {"content of a", "content of b"}
+
+
+def test_process_sdk_messages_cache_tokens():
+    messages = [
+        UserMessage(content="Hello"),
+        AssistantMessage(
+            content=[TextBlock(text="Hi!")],
+            model="claude-sonnet-4-20250514",
+        ),
+        ResultMessage(
+            subtype="success",
+            duration_ms=5000,
+            duration_api_ms=4000,
+            is_error=False,
+            num_turns=1,
+            session_id="cache-session",
+            usage={
+                "input_tokens": 36,
+                "cache_creation_input_tokens": 23554,
+                "cache_read_input_tokens": 139035,
+                "output_tokens": 3344,
+            },
+        ),
+    ]
+
+    trace = process_sdk_messages(messages, "cache-session")
+
+    assert trace is not None
+    root_span = trace.data.spans[0]
+
+    # input_tokens is the non-cached input the Anthropic API reports, matching
+    # mlflow.anthropic.autolog. Cache fields are exposed as separate keys so
+    # consumers can compute cache hit rate without scraping transcripts.
+    token_usage = root_span.get_attribute(SpanAttributeKey.CHAT_USAGE)
+    assert token_usage["input_tokens"] == 36
+    assert token_usage["output_tokens"] == 3344
+    assert token_usage["total_tokens"] == 36 + 3344
+    assert token_usage["cache_read_input_tokens"] == 139035
+    assert token_usage["cache_creation_input_tokens"] == 23554
+
+    # Trace-level aggregation should match
+    assert trace.info.token_usage["input_tokens"] == 36
+    assert trace.info.token_usage["output_tokens"] == 3344
+
+
+# ============================================================================
+# FIND LAST USER MESSAGE INDEX TESTS
+# ============================================================================
+
+
+def test_find_last_user_message_skips_skill_injection():
+    transcript = [
+        {"type": "queue-operation"},
+        {"type": "queue-operation"},
+        # Entry 2: actual user prompt
+        {
+            "type": "user",
+            "message": {"role": "user", "content": "Enable tracing on the agent."},
+            "timestamp": "2025-01-01T00:00:00Z",
+        },
+        # Entry 3: assistant thinking
+        {
+            "type": "assistant",
+            "message": {
+                "role": "assistant",
+                "content": [{"type": "thinking", "thinking": "Let me use the skill."}],
+            },
+            "timestamp": "2025-01-01T00:00:01Z",
+        },
+        # Entry 4: assistant invokes Skill tool
+        {
+            "type": "assistant",
+            "message": {
+                "role": "assistant",
+                "content": [
+                    {
+                        "type": "tool_use",
+                        "id": "toolu_abc123",
+                        "name": "Skill",
+                        "input": {"skill": "instrumenting-with-mlflow-tracing"},
+                    }
+                ],
+            },
+            "timestamp": "2025-01-01T00:00:02Z",
+        },
+        # Entry 5: tool result with commandName (correctly skipped by toolUseResult check)
+        {
+            "type": "user",
+            "toolUseResult": {
+                "success": True,
+                "commandName": "instrumenting-with-mlflow-tracing",
+            },
+            "message": {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "tool_result",
+                        "tool_use_id": "toolu_abc123",
+                        "content": "Launching skill: instrumenting-with-mlflow-tracing",
+                    }
+                ],
+            },
+            "timestamp": "2025-01-01T00:00:03Z",
+        },
+        # Entry 6: skill content injection (BUG: not flagged as tool result)
+        {
+            "type": "user",
+            "message": {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": (
+                            "Base directory for this skill: /path/to/skill\n\n"
+                            "# MLflow Tracing Guide\n\n...(full skill content)..."
+                        ),
+                    }
+                ],
+            },
+            "timestamp": "2025-01-01T00:00:04Z",
+        },
+        # Entry 7: assistant continues
+        {
+            "type": "assistant",
+            "message": {
+                "role": "assistant",
+                "content": [{"type": "thinking", "thinking": "Now let me implement tracing."}],
+            },
+            "timestamp": "2025-01-01T00:00:05Z",
+        },
+        # Entry 8: assistant text response
+        {
+            "type": "assistant",
+            "message": {
+                "role": "assistant",
+                "content": [{"type": "text", "text": "I've enabled tracing on the agent."}],
+            },
+            "timestamp": "2025-01-01T00:00:06Z",
+        },
+    ]
+
+    idx = find_last_user_message_index(transcript)
+
+    # Should return index 2 (actual user prompt), not 6 (skill injection)
+    assert idx == 2
+    assert transcript[idx]["message"]["content"] == "Enable tracing on the agent."
+
+
+def test_find_last_user_message_index_basic():
+    transcript = [
+        {"type": "queue-operation"},
+        {
+            "type": "user",
+            "message": {"role": "user", "content": "First question"},
+            "timestamp": "2025-01-01T00:00:00Z",
+        },
+        {
+            "type": "assistant",
+            "message": {
+                "role": "assistant",
+                "content": [{"type": "text", "text": "First answer"}],
+            },
+            "timestamp": "2025-01-01T00:00:01Z",
+        },
+        {
+            "type": "user",
+            "message": {"role": "user", "content": "Second question"},
+            "timestamp": "2025-01-01T00:00:02Z",
+        },
+        {
+            "type": "assistant",
+            "message": {
+                "role": "assistant",
+                "content": [{"type": "text", "text": "Second answer"}],
+            },
+            "timestamp": "2025-01-01T00:00:03Z",
+        },
+    ]
+
+    idx = find_last_user_message_index(transcript)
+
+    assert idx == 3
+    assert transcript[idx]["message"]["content"] == "Second question"
+
+
+def test_find_last_user_message_skips_consecutive_skill_injections():
+    transcript = [
+        # Entry 0: actual user prompt
+        {
+            "type": "user",
+            "message": {"role": "user", "content": "Do the thing."},
+            "timestamp": "2025-01-01T00:00:00Z",
+        },
+        # Entry 1: assistant invokes first Skill
+        {
+            "type": "assistant",
+            "message": {
+                "role": "assistant",
+                "content": [
+                    {
+                        "type": "tool_use",
+                        "id": "toolu_1",
+                        "name": "Skill",
+                        "input": {"skill": "skill-one"},
+                    }
+                ],
+            },
+            "timestamp": "2025-01-01T00:00:01Z",
+        },
+        # Entry 2: first skill tool result
+        {
+            "type": "user",
+            "toolUseResult": {"success": True, "commandName": "skill-one"},
+            "message": {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "tool_result",
+                        "tool_use_id": "toolu_1",
+                        "content": "Launching skill: skill-one",
+                    }
+                ],
+            },
+            "timestamp": "2025-01-01T00:00:02Z",
+        },
+        # Entry 3: first skill content injection
+        {
+            "type": "user",
+            "message": {
+                "role": "user",
+                "content": [{"type": "text", "text": "Base directory: /skill-one\n# Skill One"}],
+            },
+            "timestamp": "2025-01-01T00:00:03Z",
+        },
+        # Entry 4: assistant invokes second Skill
+        {
+            "type": "assistant",
+            "message": {
+                "role": "assistant",
+                "content": [
+                    {
+                        "type": "tool_use",
+                        "id": "toolu_2",
+                        "name": "Skill",
+                        "input": {"skill": "skill-two"},
+                    }
+                ],
+            },
+            "timestamp": "2025-01-01T00:00:04Z",
+        },
+        # Entry 5: second skill tool result
+        {
+            "type": "user",
+            "toolUseResult": {"success": True, "commandName": "skill-two"},
+            "message": {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "tool_result",
+                        "tool_use_id": "toolu_2",
+                        "content": "Launching skill: skill-two",
+                    }
+                ],
+            },
+            "timestamp": "2025-01-01T00:00:05Z",
+        },
+        # Entry 6: second skill content injection
+        {
+            "type": "user",
+            "message": {
+                "role": "user",
+                "content": [{"type": "text", "text": "Base directory: /skill-two\n# Skill Two"}],
+            },
+            "timestamp": "2025-01-01T00:00:06Z",
+        },
+        # Entry 7: assistant response
+        {
+            "type": "assistant",
+            "message": {
+                "role": "assistant",
+                "content": [{"type": "text", "text": "Done."}],
+            },
+            "timestamp": "2025-01-01T00:00:07Z",
+        },
+    ]
+
+    idx = find_last_user_message_index(transcript)
+
+    # Should skip both skill injections (entries 3 and 6) and return entry 0
+    assert idx == 0
+    assert transcript[idx]["message"]["content"] == "Do the thing."
+
+
+def test_process_transcript_captures_claude_code_version(tmp_path):
+    transcript = [
+        {
+            "type": "queue-operation",
+            "operation": "dequeue",
+            "timestamp": "2025-01-15T09:59:59.000Z",
+            "sessionId": "test-version-session",
+        },
+        {
+            "type": "user",
+            "version": "2.1.34",
+            "message": {"role": "user", "content": "Hello!"},
+            "timestamp": "2025-01-15T10:00:00.000Z",
+        },
+        {
+            "type": "assistant",
+            "version": "2.1.34",
+            "message": {
+                "role": "assistant",
+                "content": [{"type": "text", "text": "Hi there!"}],
+            },
+            "timestamp": "2025-01-15T10:00:01.000Z",
+        },
+    ]
+
+    transcript_path = tmp_path / "version_transcript.jsonl"
+    transcript_path.write_text("\n".join(json.dumps(entry) for entry in transcript) + "\n")
+    trace = process_transcript(str(transcript_path), "test-version-session")
+
+    assert trace is not None
+    assert trace.info.trace_metadata.get(METADATA_KEY_CLAUDE_CODE_VERSION) == "2.1.34"
+
+
+def test_process_transcript_no_version_field(mock_transcript_file):
+    trace = process_transcript(mock_transcript_file, "test-session-no-version")
+
+    assert trace is not None
+    assert METADATA_KEY_CLAUDE_CODE_VERSION not in trace.info.trace_metadata
+
+
+def test_process_transcript_includes_steer_messages(tmp_path):
+    transcript = [
+        {
+            "type": "user",
+            "message": {"role": "user", "content": "Tell me about Python."},
+            "timestamp": "2025-01-15T10:00:00.000Z",
+        },
+        {
+            "type": "assistant",
+            "message": {
+                "role": "assistant",
+                "content": [{"type": "text", "text": "Python is a programming language."}],
+            },
+            "timestamp": "2025-01-15T10:00:01.000Z",
+        },
+        {
+            "type": "queue-operation",
+            "operation": "enqueue",
+            "content": "also tell me about Java",
+            "timestamp": "2025-01-15T10:00:02.000Z",
+            "sessionId": "test-steer-session",
+        },
+        {
+            "type": "queue-operation",
+            "operation": "remove",
+            "timestamp": "2025-01-15T10:00:03.000Z",
+            "sessionId": "test-steer-session",
+        },
+        {
+            "type": "assistant",
+            "message": {
+                "role": "assistant",
+                "content": [{"type": "text", "text": "Java is also a programming language."}],
+            },
+            "timestamp": "2025-01-15T10:00:04.000Z",
+        },
+    ]
+
+    transcript_path = tmp_path / "steer_transcript.jsonl"
+    transcript_path.write_text("\n".join(json.dumps(entry) for entry in transcript) + "\n")
+    trace = process_transcript(str(transcript_path), "test-steer-session")
+    assert trace is not None
+
+    spans = list(trace.search_spans())
+    llm_spans = [s for s in spans if s.span_type == SpanType.LLM]
+    assert len(llm_spans) == 2
+
+    # The second LLM span should include the steer message in its inputs
+    second_llm = llm_spans[1]
+    input_messages = second_llm.inputs["messages"]
+    steer_messages = [m for m in input_messages if m.get("content") == "also tell me about Java"]
+    assert len(steer_messages) == 1
+    assert steer_messages[0]["role"] == "user"