Add AI Tools plugin detection to the user-agent

cmd/root/root.go

@@ -80,6 +80,7 @@ func New(ctx context.Context) *cobra.Command {
 		ctx = withCommandExecIdInUserAgent(ctx)
 		ctx = withUpstreamInUserAgent(ctx)
 		ctx = withInteractiveModeInUserAgent(ctx)
+		ctx = withAiToolsInUserAgent(ctx)
 		ctx = InjectTestPidToUserAgent(ctx)
 		cmd.SetContext(ctx)
 

cmd/root/user_agent_aitools.go

@@ -0,0 +1,26 @@
+// This file adds one aitools/<tool>_<version> pair to the user-agent for each
+// coding tool that has the Databricks plugin installed, e.g.
+// "aitools/claude-code_0.2.9 aitools/codex_0.2.9". Nothing is added when no tool
+// has it installed. The version comes from the tool's own plugin cache, falling
+// back to the CLI install state.
+package root
+
+import (
+	"context"
+
+	"github.com/databricks/cli/libs/aitools/installer"
+	"github.com/databricks/databricks-sdk-go/useragent"
+)
+
+// Key in the user agent.
+const aiToolsKey = "aitools"
+
+func withAiToolsInUserAgent(ctx context.Context) context.Context {
+	// Each tool appends one aitools/<tool>_<version> pair to the same context;
+	// these accumulate, so this is not the accidental context nesting fatcontext
+	// flags.
+	for _, tool := range installer.InstalledTools(ctx) {
+		ctx = useragent.InContext(ctx, aiToolsKey, tool) //nolint:fatcontext
+	}
+	return ctx
+}

cmd/root/user_agent_aitools_test.go

@@ -0,0 +1,45 @@
+package root
+
+import (
+	"testing"
+
+	"github.com/databricks/cli/libs/aitools/installer"
+	"github.com/databricks/databricks-sdk-go/useragent"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// recordPlugin marks the agent's plugin installed in the global CLI state, the
+// simplest signal that drives InstalledTools for any agent.
+func recordPlugin(t *testing.T, name, version string) {
+	t.Helper()
+	dir, err := installer.GlobalSkillsDir(t.Context())
+	require.NoError(t, err)
+	require.NoError(t, installer.SaveState(dir, &installer.InstallState{
+		SchemaVersion: 2,
+		Plugins:       map[string]installer.PluginRecord{name: {Plugin: "databricks", Version: version}},
+	}))
+}
+
+func TestWithAiToolsInUserAgent(t *testing.T) {
+	t.Run("none installed emits nothing", func(t *testing.T) {
+		tmp := t.TempDir()
+		t.Setenv("HOME", tmp)
+		t.Setenv("USERPROFILE", tmp)
+		t.Chdir(t.TempDir())
+
+		ua := useragent.FromContext(withAiToolsInUserAgent(t.Context()))
+		assert.NotContains(t, ua, "aitools/")
+	})
+
+	t.Run("installed tool becomes an aitools version pair", func(t *testing.T) {
+		tmp := t.TempDir()
+		t.Setenv("HOME", tmp)
+		t.Setenv("USERPROFILE", tmp)
+		t.Chdir(t.TempDir())
+		recordPlugin(t, "codex", "0.2.9")
+
+		ua := useragent.FromContext(withAiToolsInUserAgent(t.Context()))
+		assert.Contains(t, ua, "aitools/codex_0.2.9")
+	})
+}

libs/aitools/agents/agents.go

@@ -47,6 +47,12 @@ type Agent struct {
 	// Plugin describes the databricks plugin for this agent, or nil when the
 	// agent has no plugin and skills files are its native delivery.
 	Plugin *PluginSpec
+	// pluginVersion reads the installed databricks plugin version from the
+	// agent's own plugin manifest. Each agent's manifest format differs, so it is
+	// set per agent (only for formats we have verified); when nil,
+	// DatabricksPluginVersion reports no version. New agents extend support by
+	// providing their own reader here.
+	pluginVersion func(ctx context.Context, a *Agent) (string, bool)
 }
 
 // Detected returns true if the agent is installed on the system.
@@ -155,6 +161,7 @@ var Registry = []*Agent{
 		ProjectConfigDir:     ".claude",
 		Binary:               "claude",
 		Plugin:               claudePlugin(),
+		pluginVersion:        claudePluginVersion,
 	},
 	{
 		Name:        NameCursor,

libs/aitools/agents/plugins.go

@@ -6,6 +6,8 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+
+	"golang.org/x/mod/semver"
 )
 
 // pluginManifestFile is the file an agent's own plugin CLI writes to record
@@ -21,16 +23,28 @@ type pluginInstall struct {
 }
 
 // DatabricksPluginVersion reports the recorded version of the databricks plugin
-// in the agent's own plugin manifest (<ConfigDir>/plugins/installed_plugins.json,
-// Claude Code's format) and whether the plugin is installed at all. It matches
-// an "<id>@<marketplace>" key whose id is the databricks plugin. This catches
-// installs done through `databricks aitools install` and a direct `<agent>
-// plugin install`, unlike the CLI's own state file, which only records
-// CLI-driven installs. Any read/parse failure reports ("", false), so an
-// unreadable manifest reads as "not installed". When the plugin is recorded for
-// more than one scope, the first recorded version is returned (they match in the
-// common single-scope case); the version is "" when installed but unversioned.
+// in the agent's own plugin manifest and whether the plugin is installed at all.
+// This catches installs done through `databricks aitools install` and a direct
+// `<agent> plugin install`, unlike the CLI's own state file, which only records
+// CLI-driven installs. Each agent's manifest format differs, so parsing is
+// delegated to a per-agent reader; agents without a verified reader report
+// ("", false), i.e. "not installed as far as we can tell".
 func (a *Agent) DatabricksPluginVersion(ctx context.Context) (string, bool) {
+	if a.pluginVersion == nil {
+		return "", false
+	}
+	return a.pluginVersion(ctx, a)
+}
+
+// claudePluginVersion reads the databricks plugin version from Claude Code's
+// installed_plugins.json (<ConfigDir>/plugins/), which keys each entry as
+// "<id>@<marketplace>" (e.g. "databricks@claude-plugins-official") with one
+// record per install scope. Any read/parse failure reports ("", false), so an
+// unreadable manifest reads as "not installed". When the plugin is recorded for
+// more than one scope, the highest recorded version is returned (they match in
+// the common single-scope case, but the cache can retain a stale scope); the
+// version is "" when installed but unversioned.
+func claudePluginVersion(ctx context.Context, a *Agent) (string, bool) {
 	configDir, err := a.ConfigDir(ctx)
 	if err != nil {
 		return "", false
@@ -50,10 +64,15 @@ func (a *Agent) DatabricksPluginVersion(ctx context.Context) (string, bool) {
 		if id != databricksPluginID {
 			continue
 		}
-		if len(installs) > 0 {
-			return installs[0].Version, true
+		// Installed. Report the highest version across scopes;
+		best := ""
+		for _, install := range installs {
+			// Compare as semver (versions are unprefixed, e.g. "0.2.9").
+			if install.Version != "" && (best == "" || semver.Compare("v"+install.Version, "v"+best) > 0) {
+				best = install.Version
+			}
 		}
-		return "", true
+		return best, true
 	}
 	return "", false
 }

libs/aitools/agents/plugins_test.go

@@ -10,12 +10,14 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-// agentWithConfigDir returns an Agent whose ConfigDir resolves to dir.
+// agentWithConfigDir returns an Agent whose ConfigDir resolves to dir, using
+// Claude Code's manifest reader (the format these tests write).
 func agentWithConfigDir(dir string) *Agent {
 	return &Agent{
-		Name:        "test-agent",
-		DisplayName: "Test Agent",
-		ConfigDir:   func(_ context.Context) (string, error) { return dir, nil },
+		Name:          "test-agent",
+		DisplayName:   "Test Agent",
+		ConfigDir:     func(_ context.Context) (string, error) { return dir, nil },
+		pluginVersion: claudePluginVersion,
 	}
 }
 
@@ -37,6 +39,7 @@ func TestDatabricksPluginVersion(t *testing.T) {
 		{"databricks plugin present", `{"plugins":{"databricks@claude-plugins-official":[{"scope":"user","version":"0.2.6"}]}}`, true, "0.2.6"},
 		{"other plugin only", `{"plugins":{"gopls-lsp@claude-plugins-official":[{"version":"1.0.0"}]}}`, false, ""},
 		{"databricks among others", `{"plugins":{"gopls-lsp@m":[{"version":"1.0.0"}],"databricks@claude-plugins-official":[{"version":"0.3.0"}]}}`, true, "0.3.0"},
+		{"highest version across scopes", `{"plugins":{"databricks@m":[{"scope":"user","version":"0.2.9"},{"scope":"project","version":"0.2.10"},{"scope":"local","version":"0.2.8"}]}}`, true, "0.2.10"},
 		{"databricks-prefixed but distinct id does not match", `{"plugins":{"databricks-foo@m":[{"version":"9.9.9"}]}}`, false, ""},
 		{"installed but no records", `{"plugins":{"databricks@m":[]}}`, true, ""},
 		{"installed but unversioned record", `{"plugins":{"databricks@m":[{"scope":"user"}]}}`, true, ""},
@@ -55,3 +58,18 @@ func TestDatabricksPluginVersion(t *testing.T) {
 		})
 	}
 }
+
+func TestDatabricksPluginVersionWithoutReader(t *testing.T) {
+	// An agent with no per-agent manifest reader (its format is unverified)
+	// reports no version, even when an installed_plugins.json is present, so its
+	// manifest is never parsed with another agent's format.
+	configDir := t.TempDir()
+	writeManifest(t, configDir, `{"plugins":{"databricks@m":[{"version":"0.2.9"}]}}`)
+	a := &Agent{
+		Name:      "codex",
+		ConfigDir: func(_ context.Context) (string, error) { return configDir, nil },
+	}
+	version, ok := a.DatabricksPluginVersion(t.Context())
+	assert.False(t, ok)
+	assert.Empty(t, version)
+}

libs/aitools/installer/detect.go

@@ -0,0 +1,64 @@
+package installer
+
+import (
+	"context"
+	"slices"
+
+	"github.com/databricks/cli/libs/aitools/agents"
+	"github.com/databricks/cli/libs/log"
+)
+
+// InstalledTools returns "<agent>_<version>" tokens for every coding tool that
+// has the Databricks plugin installed, sorted. The version is taken from the
+// agent's own plugin cache (ground truth, catches installs made directly through
+// the agent), falling back to the version recorded in the CLI install state when
+// the plugin is not found on disk. An agent with no resolvable version is
+// omitted rather than reported without one.
+func InstalledTools(ctx context.Context) []string {
+	recorded := recordedPluginVersions(ctx)
+
+	var tools []string
+	for _, a := range agents.Registry {
+		if a.Plugin == nil {
+			continue
+		}
+		version, _ := a.DatabricksPluginVersion(ctx)
+		if version == "" {
+			version = recorded[a.Name]
+		}
+		if version != "" {
+			tools = append(tools, a.Name+"_"+version)
+		}
+	}
+	slices.Sort(tools)
+	return tools
+}
+
+// recordedPluginVersions maps agent name to the databricks plugin version
+// recorded in the CLI install state, merged across global and project scope.
+// When both scopes record a version, the project scope wins (it is the more
+// specific, closer-to-the-invocation install).
+func recordedPluginVersions(ctx context.Context) map[string]string {
+	versions := map[string]string{}
+	// Load global first, then project, so project overwrites on conflict.
+	for _, scope := range []string{ScopeGlobal, ScopeProject} {
+		dir, err := skillsDir(ctx, scope)
+		if err != nil {
+			continue
+		}
+		state, err := LoadState(dir)
+		if err != nil {
+			log.Debugf(ctx, "Could not load AI Tools install state in %s: %v", dir, err)
+			continue
+		}
+		if state == nil {
+			continue
+		}
+		for name, rec := range state.Plugins {
+			if rec.Version != "" {
+				versions[name] = rec.Version
+			}
+		}
+	}
+	return versions
+}

libs/aitools/installer/detect_test.go

@@ -0,0 +1,112 @@
+package installer
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestInstalledTools exercises the real agent registry: claude-code reads its
+// authoritative manifest (installed_plugins.json), while agents without a
+// verified manifest reader (e.g. codex, copilot) fall back to the CLI install
+// state. HOME is redirected to a temp dir so both signals resolve there.
+func TestInstalledTools(t *testing.T) {
+	tests := []struct {
+		name string
+		// claudeManifestVersion, when set, is written to claude-code's manifest.
+		claudeManifestVersion string
+		// globalPlugins / projectPlugins map agent -> recorded state version.
+		globalPlugins  map[string]string
+		projectPlugins map[string]string
+		want           []string
+	}{
+		{
+			name: "none installed",
+			want: nil,
+		},
+		{
+			// Direct-from-marketplace install: claude-code version from manifest.
+			name:                  "claude manifest version",
+			claudeManifestVersion: "0.2.9",
+			want:                  []string{"claude-code_0.2.9"},
+		},
+		{
+			// codex has no verified manifest reader, so it uses the state version.
+			name:          "state fallback for codex",
+			globalPlugins: map[string]string{"codex": "0.2.7"},
+			want:          []string{"codex_0.2.7"},
+		},
+		{
+			// Manifest wins over state for claude-code when both are present.
+			name:                  "manifest overrides state",
+			claudeManifestVersion: "0.2.9",
+			globalPlugins:         map[string]string{"claude-code": "0.2.5"},
+			want:                  []string{"claude-code_0.2.9"},
+		},
+		{
+			// Project scope wins over global in the state fallback.
+			name:           "project state overrides global",
+			globalPlugins:  map[string]string{"codex": "0.2.5"},
+			projectPlugins: map[string]string{"codex": "0.2.7"},
+			want:           []string{"codex_0.2.7"},
+		},
+		{
+			// Manifest (claude) and state (codex) contribute; output is sorted.
+			name:                  "mixed sources sorted",
+			claudeManifestVersion: "0.2.9",
+			globalPlugins:         map[string]string{"codex": "0.2.8"},
+			want:                  []string{"claude-code_0.2.9", "codex_0.2.8"},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			tmp := t.TempDir()
+			t.Setenv("HOME", tmp)
+			t.Setenv("USERPROFILE", tmp)
+			// Fresh cwd so project-scope state does not leak from the package dir.
+			t.Chdir(t.TempDir())
+
+			if tc.claudeManifestVersion != "" {
+				writeClaudeManifest(t, filepath.Join(tmp, ".claude"), tc.claudeManifestVersion)
+			}
+			savePluginState(t, GlobalSkillsDir, tc.globalPlugins)
+			savePluginState(t, ProjectSkillsDir, tc.projectPlugins)
+
+			assert.Equal(t, tc.want, InstalledTools(t.Context()))
+		})
+	}
+}
+
+// writeClaudeManifest writes a Claude installed_plugins.json recording the
+// databricks plugin at the given version under configDir.
+func writeClaudeManifest(t *testing.T, configDir, version string) {
+	t.Helper()
+	require.NoError(t, os.MkdirAll(filepath.Join(configDir, "plugins"), 0o755))
+	body := `{"version":2,"plugins":{"databricks@claude-plugins-official":[{"version":"` + version + `"}]}}`
+	require.NoError(t, os.WriteFile(filepath.Join(configDir, "plugins", "installed_plugins.json"), []byte(body), 0o644))
+}
+
+// savePluginState writes an install state recording a databricks plugin version
+// for each agent into the directory returned by dirFn. It is a no-op when the
+// map is empty, so a scope with no plugin installs leaves no state file.
+func savePluginState(t *testing.T, dirFn func(context.Context) (string, error), versions map[string]string) {
+	t.Helper()
+	if len(versions) == 0 {
+		return
+	}
+	plugins := make(map[string]PluginRecord, len(versions))
+	for name, v := range versions {
+		plugins[name] = PluginRecord{Plugin: "databricks", Version: v}
+	}
+	dir, err := dirFn(t.Context())
+	require.NoError(t, err)
+	require.NoError(t, SaveState(dir, &InstallState{
+		SchemaVersion: schemaVersionV2,
+		Plugins:       plugins,
+	}))
+}