chore: import upstream snapshot with attribution

mlflow/claude_code/cli.py

@@ -0,0 +1,352 @@
+"""MLflow CLI commands for Claude Code integration."""
+
+import os
+import sys
+from pathlib import Path
+
+import click
+
+from mlflow.claude_code.config import get_tracing_status, setup_environment_config
+from mlflow.claude_code.hooks import stop_hook_handler
+from mlflow.claude_code.plugin import (
+    disable_tracing_plugin,
+    ensure_plugin_installed,
+)
+from mlflow.environment_variables import (
+    MLFLOW_EXPERIMENT_ID,
+    MLFLOW_EXPERIMENT_NAME,
+    MLFLOW_TRACKING_URI,
+)
+
+
+def _title(text: str) -> str:
+    return click.style(text, fg="magenta", bold=True)
+
+
+def _ok(text: str) -> str:
+    return click.style(text, fg="green", bold=True)
+
+
+def _warn(text: str) -> str:
+    return click.style(text, fg="yellow", bold=True)
+
+
+def _error(text: str) -> str:
+    return click.style(text, fg="red", bold=True)
+
+
+def _label(text: str) -> str:
+    return click.style(text, bold=True)
+
+
+def _question(text: str) -> str:
+    return click.style(text, fg="yellow", bold=True)
+
+
+def _value(text: str) -> str:
+    return click.style(text, fg="cyan")
+
+
+def _muted(text: str) -> str:
+    return click.style(text, dim=True)
+
+
+_DEFAULT_TRACKING_URI_SENTINEL = "default"
+
+
+@click.group("autolog")
+def commands():
+    """Commands for autologging with MLflow."""
+
+
+@commands.group("claude", invoke_without_command=True)
+@click.option(
+    "--directory",
+    "-d",
+    default=".",
+    type=click.Path(file_okay=False, dir_okay=True),
+    help="Directory to set up tracing in (default: current directory)",
+)
+@click.option(
+    "--tracking-uri", "-u", help="MLflow tracking URI (e.g., 'databricks' or 'file://mlruns')"
+)
+@click.option("--experiment-id", "-e", help="MLflow experiment ID")
+@click.option("--experiment-name", "-n", help="MLflow experiment name")
+@click.option(
+    "--disable",
+    is_flag=True,
+    help="Disable Claude tracing (removes config from both settings.json and settings.local.json)",
+)
+@click.option("--status", is_flag=True, help="Show current tracing status")
+@click.option(
+    "--local",
+    is_flag=True,
+    help="Write config to settings.local.json instead of settings.json during setup.",
+)
+@click.option(
+    "--non-interactive",
+    "-y",
+    is_flag=True,
+    help="Skip prompts and use flags, environment variables, or defaults.",
+)
+@click.option(
+    "--mlflow-cmd",
+    default=None,
+    help=(
+        "Deprecated and ignored. Python-based Claude hooks were replaced by the "
+        "marketplace plugin runtime."
+    ),
+)
+@click.pass_context
+def claude(
+    ctx: click.Context,
+    directory: str,
+    tracking_uri: str | None,
+    experiment_id: str | None,
+    experiment_name: str | None,
+    disable: bool,
+    status: bool,
+    local: bool,
+    non_interactive: bool,
+    mlflow_cmd: str | None,
+) -> None:
+    """Set up Claude Code tracing in a directory.
+
+    This command installs the MLflow Claude plugin into Claude Code and writes
+    MLflow configuration into `.claude/settings.json`. After setup, use the
+    regular `claude` command and traces will be created by the plugin runtime.
+
+    Examples:
+
+      # Set up tracing in current directory with local storage
+      mlflow autolog claude
+
+      # Set up tracing in a specific project directory
+      mlflow autolog claude -d ~/my-project
+
+      # Set up tracing with Databricks
+      mlflow autolog claude -u databricks -e 123456789
+
+      # Set up tracing with custom tracking URI
+      mlflow autolog claude -u file://./custom-mlruns
+
+      # Disable tracing in current directory
+      mlflow autolog claude --disable
+    """
+    # Skip setup when a subcommand (e.g., stop-hook) is being invoked
+    if ctx.invoked_subcommand is not None:
+        return
+
+    if experiment_id and experiment_name:
+        raise click.BadParameter("Choose either --experiment-id or --experiment-name, not both.")
+
+    if mlflow_cmd is not None:
+        if not mlflow_cmd.strip():
+            raise click.BadParameter(
+                "must not be empty or whitespace-only", param_hint="'--mlflow-cmd'"
+            )
+        click.echo(f"{_warn('⚠')} {_muted('--mlflow-cmd is deprecated and ignored.')}")
+
+    if local and (status or disable):
+        raise click.UsageError(
+            "--local can only be used during setup, not with --status or --disable"
+        )
+
+    target_dir = Path(directory).resolve()
+    claude_dir = target_dir / ".claude"
+    settings_file = claude_dir / "settings.json"
+    local_settings_file = claude_dir / "settings.local.json"
+
+    if status:
+        _show_status(target_dir, settings_file)
+        return
+
+    if disable:
+        removed_shared = _handle_disable(settings_file)
+        removed_local = _handle_disable(local_settings_file)
+        if not removed_shared and not removed_local:
+            click.echo(f"{_error('✗')} No Claude configuration found - tracing was not enabled")
+        return
+
+    if local:
+        settings_file = local_settings_file
+
+    _print_setup_intro(tracking_uri, experiment_id, experiment_name, non_interactive)
+    tracking_uri, experiment_id, experiment_name = _resolve_setup_inputs(
+        tracking_uri,
+        experiment_id,
+        experiment_name,
+        non_interactive,
+    )
+
+    click.echo(f"{_title('MLflow Claude Tracing Setup')}")
+    click.echo(f"{_label('Project:')} {_value(str(target_dir))}")
+
+    # Create .claude directory and install the plugin runtime
+    claude_dir.mkdir(parents=True, exist_ok=True)
+    click.echo(f"{_label('Installing plugin:')} {_muted('MLflow Claude plugin for Claude Code')}")
+    try:
+        ensure_plugin_installed(target_dir)
+    except click.ClickException:
+        raise
+    except Exception as exc:
+        raise click.ClickException(f"Failed to configure Claude tracing: {exc}") from exc
+    click.echo(f"{_ok('✓')} Claude Code plugin installed")
+
+    # Set up environment variables consumed by the plugin
+    setup_environment_config(settings_file, tracking_uri, experiment_id, experiment_name)
+
+    # Show final status
+    _show_setup_status(target_dir, settings_file)
+
+
+def _print_setup_intro(
+    tracking_uri: str | None,
+    experiment_id: str | None,
+    experiment_name: str | None,
+    non_interactive: bool,
+) -> None:
+    if non_interactive or not _is_interactive_shell():
+        return
+
+    missing_tracking = not (tracking_uri or os.environ.get(MLFLOW_TRACKING_URI.name))
+    missing_experiment = not (
+        experiment_id
+        or experiment_name
+        or os.environ.get(MLFLOW_EXPERIMENT_ID.name)
+        or os.environ.get(MLFLOW_EXPERIMENT_NAME.name)
+    )
+    if not missing_tracking and not missing_experiment:
+        return
+
+    click.echo(f"{_title('Interactive Mode')}")
+    click.echo(_muted("MLflow Claude tracing setup is running in interactive mode."))
+    click.echo(
+        _muted(
+            "If you want non-interactive setup, provide values with CLI options or set "
+            "MLFLOW_TRACKING_URI and MLFLOW_EXPERIMENT_ID in your environment."
+        )
+    )
+    click.echo("")
+
+
+def _resolve_setup_inputs(
+    tracking_uri: str | None,
+    experiment_id: str | None,
+    experiment_name: str | None,
+    non_interactive: bool,
+) -> tuple[str | None, str | None, str | None]:
+    resolved_tracking_uri = tracking_uri or os.environ.get(MLFLOW_TRACKING_URI.name)
+    resolved_experiment_id = experiment_id or os.environ.get(MLFLOW_EXPERIMENT_ID.name)
+    resolved_experiment_name = experiment_name or os.environ.get(MLFLOW_EXPERIMENT_NAME.name)
+
+    if non_interactive or not _is_interactive_shell():
+        return resolved_tracking_uri, resolved_experiment_id, resolved_experiment_name
+
+    if not resolved_tracking_uri:
+        import mlflow
+
+        actual_default_tracking_uri = mlflow.get_tracking_uri()
+        resolved_tracking_uri = click.prompt(
+            _question("MLflow tracking URI"),
+            default=_DEFAULT_TRACKING_URI_SENTINEL,
+            show_default=True,
+        ).strip()
+        if resolved_tracking_uri == _DEFAULT_TRACKING_URI_SENTINEL:
+            resolved_tracking_uri = actual_default_tracking_uri
+
+    if not resolved_experiment_id and not resolved_experiment_name:
+        resolved_experiment_id = click.prompt(
+            _question("MLflow experiment ID"),
+            default="0",
+            show_default=True,
+        ).strip()
+
+    return resolved_tracking_uri, resolved_experiment_id, resolved_experiment_name
+
+
+def _is_interactive_shell() -> bool:
+    return sys.stdin.isatty() and sys.stdout.isatty()
+
+
+def _handle_disable(settings_file: Path) -> bool:
+    """Handle disable for a single settings file.
+
+    Returns:
+        True if config was removed, False if no config found
+    """
+    if disable_tracing_plugin(settings_file):
+        click.echo(f"{_ok('✓')} Claude tracing disabled in {settings_file.name}")
+        return True
+    return False
+
+
+def _show_status(target_dir: Path, settings_file: Path) -> None:
+    """Show current tracing status."""
+    click.echo(f"{_title('MLflow Claude Tracing Status')}")
+    click.echo(f"{_label('Project:')} {_value(str(target_dir))}")
+
+    status = get_tracing_status(settings_file)
+
+    if not status.enabled:
+        click.echo(f"{_error('✗')} Claude tracing is not enabled")
+        if status.reason:
+            click.echo(f"  {_label('Reason:')} {_muted(status.reason)}")
+        return
+
+    click.echo(f"{_ok('✓')} Claude tracing is enabled")
+    click.echo(f"{_label('Tracking URI:')} {_value(str(status.tracking_uri))}")
+
+    if status.experiment_name:
+        click.echo(f"{_label('Experiment name:')} {_value(status.experiment_name)}")
+    if status.experiment_id:
+        click.echo(f"{_label('Experiment ID:')} {_value(status.experiment_id)}")
+    elif not status.experiment_name:
+        click.echo(f"{_label('Experiment:')} {_muted('Default (experiment 0)')}")
+
+
+def _show_setup_status(
+    target_dir: Path,
+    settings_file: Path,
+) -> None:
+    """Show setup completion status."""
+    current_dir = Path.cwd().resolve()
+    status = get_tracing_status(settings_file)
+
+    click.echo("")
+    click.echo(_title("Setup Complete"))
+    click.echo(f"{_label('Project:')} {_value(str(target_dir))}")
+
+    # Show tracking configuration
+    if status.tracking_uri:
+        click.echo(f"{_label('Tracking URI:')} {_value(status.tracking_uri)}")
+
+    if status.experiment_name:
+        click.echo(f"{_label('Experiment name:')} {_value(status.experiment_name)}")
+    if status.experiment_id:
+        click.echo(f"{_label('Experiment ID:')} {_value(status.experiment_id)}")
+    elif not status.experiment_name:
+        click.echo(f"{_label('Experiment:')} {_muted('Default (experiment 0)')}")
+
+    # Show next steps
+    click.echo("")
+    click.echo(_title("Next Steps"))
+
+    # Only show cd if it's a different directory
+    if target_dir != current_dir:
+        click.echo(f"  {_muted('Work from:')} {_value(str(target_dir))}")
+
+    click.echo(f"  {_muted('1.')} Use Claude Code as usual in this directory.")
+    click.echo(
+        f"  {_muted('2.')} Visit the MLflow UI after a Claude conversation ends to inspect traces."
+    )
+
+    click.echo("")
+    click.echo(_title("Disable Later"))
+    click.echo(f"  {_value('mlflow autolog claude --disable')}")
+
+
+@claude.command("stop-hook", hidden=True)
+def stop_hook() -> None:
+    """Legacy hook shim kept for older Python-hook installations."""
+    stop_hook_handler()