databricks
diff --git a/‎.github/workflows/ci.yml‎
Lines changed: 8 additions & 1 deletion b/‎.github/workflows/ci.yml‎
Lines changed: 8 additions & 1 deletion
diff --git a/‎pyproject.toml‎
Lines changed: 6 additions & 0 deletions b/‎pyproject.toml‎
Lines changed: 6 additions & 0 deletions
diff --git a/‎src/ucode/agents/claude.py‎
Lines changed: 134 additions & 1 deletion b/‎src/ucode/agents/claude.py‎
Lines changed: 134 additions & 1 deletion
diff --git a/‎src/ucode/agents/codex.py‎
Lines changed: 53 additions & 0 deletions b/‎src/ucode/agents/codex.py‎
Lines changed: 53 additions & 0 deletions
diff --git a/‎src/ucode/agents/opencode.py‎
Lines changed: 26 additions & 3 deletions b/‎src/ucode/agents/opencode.py‎
Lines changed: 26 additions & 3 deletions
@@ -15,7 +15,7 @@ jobs:
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
       - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b  # v8.1.0
-      - run: uv run pytest --ignore=tests/test_e2e.py
+      - run: uv run pytest --ignore=tests/test_e2e.py --ignore=tests/test_e2e_tracing.py
 
   e2e:
     if: vars.E2E_ENABLED == 'true'
@@ -54,3 +54,10 @@ jobs:
       # set, the auth code path doesn't shell out at all — this is a safety
       # net for any code path we may have missed.
       - run: uv run pytest tests/test_e2e.py -v < /dev/null
+      # MLflow tracing e2e lives in its own file and needs the `tracing`
+      # extra so `import mlflow` resolves (otherwise the test importorskips
+      # and silently passes as skipped). `!cancelled()` lets this run even
+      # when the previous pytest step failed — the two suites are
+      # independent and one shouldn't mask the other.
+      - if: ${{ !cancelled() }}
+        run: uv run --extra tracing pytest tests/test_e2e_tracing.py -v < /dev/null
@@ -15,6 +15,12 @@ dependencies = [
   "typer>=0.12.0",
 ]
 
+[project.optional-dependencies]
+# Only the Claude Code tracing path needs the Python MLflow runtime; OpenCode
+# and Codex trace via JS plugins. Kept out of the core deps so a plain `ucode`
+# install stays lean. Enable with `uv tool install "ucode[tracing]"`.
+tracing = ["mlflow[databricks]>=3.4"]
+
 [project.scripts]
 ucode = "ucode.cli:main"
 
 
@@ -5,6 +5,7 @@
 import os
 import re
 import shutil
+import subprocess
 from pathlib import Path
 
 from ucode.agent_updates import available_npm_package_update
@@ -23,7 +24,8 @@
 )
 from ucode.state import mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_version
-from ucode.ui import print_warning
+from ucode.tracing import tracing_env
+from ucode.ui import print_note, print_success, print_warning
 
 CLAUDE_CONFIG_DIR = Path.home() / ".claude"
 CLAUDE_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "ucode-settings.json"
@@ -61,6 +63,20 @@ def _resolve_web_search_model(state: dict) -> str | None:
 WEB_SEARCH_MCP_NAME = "web_search"
 _CLAUDE_MODEL_RE = re.compile(r"^databricks-claude-(opus|sonnet)-(\d+)-(\d+)(.*)$")
 
+# Env keys consumed by the MLflow Claude tracing plugin. Written into the
+# settings `env` block; the plugin runtime (installed separately) reads them.
+CLAUDE_TRACING_ENV_KEYS = (
+    "MLFLOW_CLAUDE_TRACING_ENABLED",
+    "MLFLOW_TRACKING_URI",
+    "MLFLOW_EXPERIMENT_ID",
+)
+CLAUDE_TRACING_MARKETPLACE = "mlflow/mlflow"
+CLAUDE_TRACING_PLUGIN = "mlflow-tracing@mlflow-plugins"
+# The plugin runtime shells out to the `mlflow` CLI, so it must be on PATH at
+# this minimum version. ucode installs/upgrades it via `uv tool`.
+MLFLOW_CLI_SPEC = "mlflow[databricks]>=3.4"
+MINIMUM_MLFLOW_VERSION = (3, 4)
+
 
 def _web_search_mcp_entry(workspace: str, search_model: str, profile: str | None = None) -> dict:
     """Stdio MCP server entry pointing at `ucode mcp web-search`. Resolves
@@ -198,8 +214,19 @@ def write_tool_config(state: dict, model: str) -> dict:
         disable_web_search=web_search_model is not None,
         profile=state.get("profile"),
     )
+    tracing_env_vars = tracing_env(state, "claude")
+    if tracing_env_vars:
+        overlay["env"]["MLFLOW_CLAUDE_TRACING_ENABLED"] = "true"
+        overlay["env"].update(tracing_env_vars)
+        managed_keys = managed_keys + [["env", key] for key in CLAUDE_TRACING_ENV_KEYS]
+
     existing = read_json_safe(CLAUDE_SETTINGS_PATH)
     merged = deep_merge_dict(existing, overlay)
+    if not tracing_env_vars:
+        env_block = merged.get("env")
+        if isinstance(env_block, dict):
+            for key in CLAUDE_TRACING_ENV_KEYS:
+                env_block.pop(key, None)
     write_json_file(CLAUDE_SETTINGS_PATH, merged)
 
     if web_search_model:
@@ -210,6 +237,112 @@ def write_tool_config(state: dict, model: str) -> dict:
     return state
 
 
+def ensure_tracing_runtime() -> bool:
+    """Ensure Claude's MLflow tracing runtime is ready: an `mlflow` CLI >= 3.4 on
+    PATH (the plugin shells out to it) and the MLflow Claude plugin installed.
+
+    Best-effort — warns and returns False if a piece can't be set up, so
+    `ucode configure tracing` can still finish for other agents."""
+    if not _ensure_mlflow_cli():
+        return False
+    return _install_claude_tracing_plugin()
+
+
+def _parse_mlflow_version(text: str) -> tuple[int, int] | None:
+    match = re.search(r"(\d+)\.(\d+)", text)
+    if not match:
+        return None
+    return int(match.group(1)), int(match.group(2))
+
+
+def _installed_mlflow_version() -> tuple[int, int] | None:
+    """The (major, minor) of the `mlflow` CLI on PATH, or None if absent."""
+    if not shutil.which("mlflow"):
+        return None
+    try:
+        result = subprocess.run(
+            ["mlflow", "--version"], check=False, capture_output=True, text=True, timeout=30
+        )
+    except (OSError, subprocess.TimeoutExpired):
+        return None
+    return _parse_mlflow_version(result.stdout or result.stderr or "")
+
+
+def _ensure_mlflow_cli() -> bool:
+    """Ensure an `mlflow` CLI >= 3.4 is on PATH, installing or upgrading it via
+    `uv tool` when needed."""
+    current = _installed_mlflow_version()
+    if current and current >= MINIMUM_MLFLOW_VERSION:
+        return True
+
+    if not shutil.which("uv"):
+        verb = "upgrade" if current else "install"
+        print_warning(
+            f"Claude tracing needs the `mlflow` CLI >= 3.4 on PATH, but `uv` is not "
+            f'available to {verb} it. Run `uv tool install "{MLFLOW_CLI_SPEC}"` '
+            f'(or `pip install "{MLFLOW_CLI_SPEC}"`), then re-run `ucode configure tracing`.'
+        )
+        return False
+
+    print_note(f"{'Upgrading' if current else 'Installing'} the mlflow CLI ({MLFLOW_CLI_SPEC})...")
+    # --force replaces an existing (older) uv-managed mlflow tool in place.
+    cmd = ["uv", "tool", "install", MLFLOW_CLI_SPEC]
+    if current:
+        cmd.append("--force")
+    try:
+        subprocess.run(cmd, check=True, timeout=600)
+    except (OSError, subprocess.CalledProcessError, subprocess.TimeoutExpired) as exc:
+        print_warning(f"Could not install the mlflow CLI automatically: {exc}")
+        return False
+
+    if not shutil.which("mlflow"):
+        print_warning(
+            "Installed mlflow, but `mlflow` is still not on PATH. Ensure your uv tool "
+            "bin directory (e.g. ~/.local/bin) is on PATH, then re-run `ucode configure tracing`."
+        )
+        return False
+    print_success("mlflow CLI ready")
+    return True
+
+
+def _install_claude_tracing_plugin() -> bool:
+    binary = SPEC["binary"]
+    if not shutil.which(binary):
+        print_warning("`claude` is not installed; skipping MLflow tracing plugin install.")
+        return False
+    commands = [
+        [
+            binary,
+            "plugin",
+            "marketplace",
+            "add",
+            CLAUDE_TRACING_MARKETPLACE,
+            "--sparse",
+            ".claude-plugin",
+        ],
+        [binary, "plugin", "install", CLAUDE_TRACING_PLUGIN],
+    ]
+    for cmd in commands:
+        try:
+            result = subprocess.run(cmd, check=False, capture_output=True, text=True, timeout=120)
+        except (OSError, subprocess.TimeoutExpired) as exc:
+            print_warning(f"Could not install the Claude MLflow plugin: {exc}")
+            return False
+        if result.returncode != 0:
+            output = (result.stderr or result.stdout or "").strip()
+            last = output.splitlines()[-1] if output else f"exit {result.returncode}"
+            # `marketplace add` / `install` are idempotent; treat "already
+            # added/installed" as success and keep going. Best-effort match
+            # against stderr — an upstream wording change would degrade this
+            # to a noisy warning on re-runs, but never corrupts state.
+            if "already" in last.lower():
+                continue
+            print_warning(f"Claude MLflow plugin step failed: {last}")
+            return False
+    print_success("Claude MLflow tracing plugin installed")
+    return True
+
+
 def default_model(state: dict) -> str | None:
     claude_models = state.get("claude_models") or {}
     return claude_models.get("opus") or claude_models.get("sonnet") or claude_models.get("haiku")
 
@@ -4,6 +4,7 @@
 
 import os
 import re
+import subprocess
 from pathlib import Path
 
 from ucode.agent_updates import available_npm_package_update
@@ -22,6 +23,8 @@
 )
 from ucode.state import mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_version
+from ucode.tracing import agent_tracing, apply_tracing_env
+from ucode.ui import print_success, print_warning
 
 CODEX_CONFIG_DIR = Path.home() / ".codex"
 CODEX_PROFILE_NAME = "ucode"
@@ -32,6 +35,8 @@
 CODEX_MODEL_PROVIDER_NAME = "ucode-databricks"
 MINIMUM_CODEX_VERSION = (0, 134, 0)
 MINIMUM_CODEX_VERSION_TEXT = "0.134.0"
+CODEX_TRACING_NOTIFY = ["mlflow-codex", "notify-hook"]
+CODEX_TRACING_PACKAGE = "@mlflow/codex"
 
 
 SPEC: ToolSpec = {
@@ -235,6 +240,7 @@ def write_tool_config(state: dict, model: str | None = None) -> dict:
         overlay = render_legacy_overlay(workspace, chosen_model, databricks_profile)
         doc = read_toml_safe(LEGACY_CODEX_CONFIG_PATH)
         deep_merge_dict(doc, overlay)
+        _apply_tracing_notify(doc, state)
         write_toml_file(LEGACY_CODEX_CONFIG_PATH, doc)
         state = mark_tool_managed(state, "codex", LEGACY_MANAGED_KEYS)
         save_state(state)
@@ -245,12 +251,55 @@ def write_tool_config(state: dict, model: str | None = None) -> dict:
     overlay = render_overlay(workspace, chosen_model, databricks_profile)
     doc = read_toml_safe(CODEX_CONFIG_PATH)
     deep_merge_dict(doc, overlay)
+    _apply_tracing_notify(doc, state)
     write_toml_file(CODEX_CONFIG_PATH, doc)
     state = mark_tool_managed(state, "codex", MANAGED_KEYS)
     save_state(state)
     return state
 
 
+def _apply_tracing_notify(doc: dict, state: dict) -> None:
+    """Set/clear the Codex ``notify`` hook that streams session traces to MLflow.
+
+    Only ucode's own notify value is removed on disable, so a user-defined
+    ``notify`` is left intact. When enabling on top of a pre-existing user
+    ``notify``, warn before replacing — the prior value is in the backup file
+    but the user has to restore it manually."""
+    if agent_tracing(state, "codex") is not None:
+        existing = doc.get("notify")
+        if existing is not None and list(existing) != CODEX_TRACING_NOTIFY:
+            print_warning(
+                f"Codex `notify` is already set to {existing!r}; replacing it with the "
+                "MLflow tracing hook. The previous value is preserved in the Codex "
+                "config backup — restore it manually if you need both."
+            )
+        doc["notify"] = list(CODEX_TRACING_NOTIFY)
+    elif list(doc.get("notify") or []) == CODEX_TRACING_NOTIFY:
+        doc.pop("notify", None)
+
+
+def ensure_tracing_dependency() -> bool:
+    """Install the `@mlflow/codex` npm package that provides the `mlflow-codex`
+    notify-hook binary. Best-effort: warns and returns False on failure."""
+    import shutil
+
+    if shutil.which("mlflow-codex"):
+        return True
+    if not shutil.which("npm"):
+        print_warning(
+            f"`npm` is not available to install {CODEX_TRACING_PACKAGE}; "
+            "Codex tracing will be inactive until it is installed."
+        )
+        return False
+    try:
+        subprocess.run(["npm", "install", "-g", CODEX_TRACING_PACKAGE], check=True, timeout=300)
+    except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
+        print_warning(f"Could not install {CODEX_TRACING_PACKAGE}; Codex tracing will be inactive.")
+        return False
+    print_success("Codex MLflow tracing hook installed")
+    return True
+
+
 def default_model(state: dict) -> str | None:
     """Pick the newest GPT model when multiple are available.
 
@@ -279,6 +328,10 @@ def launch(state: dict, tool_args: list[str]) -> None:
     workspace = state.get("workspace")
     if workspace:
         os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile"))
+    # The notify hook subprocess Codex spawns inherits this env, so MLflow
+    # routing flows through to it without writing a separate tracing config.
+    # When tracing is off this also clears any stale outer-shell value.
+    apply_tracing_env(os.environ, state, "codex")
     os.execvp(binary, [binary, "--profile", CODEX_PROFILE_NAME, *tool_args])
 
 
 
@@ -23,12 +23,14 @@
 )
 from ucode.state import mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_version
+from ucode.tracing import agent_tracing, apply_tracing_env
 
 OPENCODE_XDG_CONFIG_HOME = APP_DIR / "opencode-xdg"
 OPENCODE_CONFIG_DIR = OPENCODE_XDG_CONFIG_HOME / "opencode"
 OPENCODE_CONFIG_PATH = OPENCODE_CONFIG_DIR / "opencode.json"
 OPENCODE_BACKUP_PATH = APP_DIR / "opencode-config.backup.json"
 OPENCODE_MCP_AUTH_HEADER_VALUE = "Bearer {env:OAUTH_TOKEN}"
+OPENCODE_TRACING_PLUGIN = "@mlflow/opencode"
 
 SPEC: ToolSpec = {
     "binary": "opencode",
@@ -150,12 +152,29 @@ def write_tool_config(
         for stale in ("databricks-anthropic", "databricks-google", "databricks-openai"):
             providers.pop(stale, None)
     merged = deep_merge_dict(existing, overlay)
+    _apply_tracing_plugin(merged, state)
     write_json_file(OPENCODE_CONFIG_PATH, merged)
     state = mark_tool_managed(state, "opencode", managed_keys)
     save_state(state)
     return state, token
 
 
+def _apply_tracing_plugin(config: dict, state: dict) -> None:
+    """Add/remove the MLflow plugin in opencode.json's top-level ``plugin`` list
+    to match the current tracing state, leaving any user plugins untouched.
+    OpenCode auto-installs listed npm plugins at startup."""
+    plugins = config.get("plugin")
+    plugins = (
+        [p for p in plugins if p != OPENCODE_TRACING_PLUGIN] if isinstance(plugins, list) else []
+    )
+    if agent_tracing(state, "opencode") is not None:
+        plugins.append(OPENCODE_TRACING_PLUGIN)
+    if plugins:
+        config["plugin"] = plugins
+    else:
+        config.pop("plugin", None)
+
+
 def build_mcp_server_entry(url: str) -> dict:
     return {
         "type": "remote",
@@ -216,17 +235,21 @@ def _refresh_forever(state: dict, stop_event: threading.Event) -> None:
             continue
 
 
-def build_runtime_env(token: str) -> dict[str, str]:
+def build_runtime_env(token: str, state: dict | None = None) -> dict[str, str]:
     env = os.environ.copy()
     env["OAUTH_TOKEN"] = token
     env["XDG_CONFIG_HOME"] = str(OPENCODE_XDG_CONFIG_HOME)
+    if state is not None:
+        # apply_tracing_env clears the MLflow vars when tracing is off, so a
+        # stale outer-shell value can't leak through.
+        apply_tracing_env(env, state, "opencode")
     return env
 
 
 def launch(state: dict, tool_args: list[str]) -> None:
     """Launch opencode with background token refresh (same pattern as Gemini)."""
     token = _refresh_token_once(state)
-    env = build_runtime_env(token)
+    env = build_runtime_env(token, state)
 
     stop_event = threading.Event()
     refresher = threading.Thread(
@@ -257,4 +280,4 @@ def validate_env(state: dict) -> dict[str, str]:
     workspace = state.get("workspace")
     if not workspace:
         raise RuntimeError("No workspace configured.")
-    return build_runtime_env(get_databricks_token(workspace, state.get("profile")))
+    return build_runtime_env(get_databricks_token(workspace, state.get("profile")), state)