Add agent detection library

acceptance/bundle/user_agent/test.toml

@@ -4,3 +4,7 @@ IncludeRequestHeaders = ["User-Agent"]
 
 [Env]
 DATABRICKS_CACHE_ENABLED = 'false'
+# Clear agent env vars to prevent them from affecting test output
+CLAUDECODE = ''
+GEMINI_CLI = ''
+CURSOR_AGENT = ''

cmd/root/root.go

@@ -12,6 +12,7 @@ import (
 	"time"
 
 	"github.com/databricks/cli/internal/build"
+	"github.com/databricks/cli/libs/agent"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/dbr"
 	"github.com/databricks/cli/libs/log"
@@ -79,6 +80,7 @@ func New(ctx context.Context) *cobra.Command {
 		ctx = withCommandInUserAgent(ctx, cmd)
 		ctx = withCommandExecIdInUserAgent(ctx)
 		ctx = withUpstreamInUserAgent(ctx)
+		ctx = withAgentInUserAgent(ctx)
 		ctx = InjectTestPidToUserAgent(ctx)
 		cmd.SetContext(ctx)
 		return nil
@@ -129,6 +131,9 @@ Stack Trace:
 	// Detect if the CLI is running on DBR and store this on the context.
 	ctx = dbr.DetectRuntime(ctx)
 
+	// Detect if the CLI is running under an agent.
+	ctx = agent.Detect(ctx)
+
 	// Set a command execution ID value in the context
 	ctx = cmdctx.GenerateExecId(ctx)
 

cmd/root/user_agent_agent.go

@@ -0,0 +1,27 @@
+// This file integrates agent detection with the user agent string.
+//
+// The actual detection logic is in libs/agent. This file simply retrieves
+// the detected agent name from the context and adds it to the user agent.
+//
+// Example user agent strings:
+//   - With Claude Code: "cli/X.Y.Z ... agent/claude-code ..."
+//   - No agent: "cli/X.Y.Z ..." (no agent tag)
+package root
+
+import (
+	"context"
+
+	"github.com/databricks/cli/libs/agent"
+	"github.com/databricks/databricks-sdk-go/useragent"
+)
+
+// Key in the user agent
+const agentKey = "agent"
+
+func withAgentInUserAgent(ctx context.Context) context.Context {
+	product := agent.Product(ctx)
+	if product == "" {
+		return ctx
+	}
+	return useragent.InContext(ctx, agentKey, product)
+}

cmd/root/user_agent_agent_test.go

@@ -0,0 +1,42 @@
+package root
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/cli/libs/agent"
+	"github.com/databricks/databricks-sdk-go/useragent"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAgentClaudeCode(t *testing.T) {
+	ctx := context.Background()
+	ctx = agent.Mock(ctx, agent.ClaudeCode)
+
+	ctx = withAgentInUserAgent(ctx)
+	assert.Contains(t, useragent.FromContext(ctx), "agent/claude-code")
+}
+
+func TestAgentGeminiCLI(t *testing.T) {
+	ctx := context.Background()
+	ctx = agent.Mock(ctx, agent.GeminiCLI)
+
+	ctx = withAgentInUserAgent(ctx)
+	assert.Contains(t, useragent.FromContext(ctx), "agent/gemini-cli")
+}
+
+func TestAgentCursor(t *testing.T) {
+	ctx := context.Background()
+	ctx = agent.Mock(ctx, agent.Cursor)
+
+	ctx = withAgentInUserAgent(ctx)
+	assert.Contains(t, useragent.FromContext(ctx), "agent/cursor")
+}
+
+func TestAgentNotSet(t *testing.T) {
+	ctx := context.Background()
+	ctx = agent.Mock(ctx, "")
+
+	ctx = withAgentInUserAgent(ctx)
+	assert.NotContains(t, useragent.FromContext(ctx), "agent/")
+}

libs/agent/agent.go

@@ -0,0 +1,76 @@
+package agent
+
+import (
+	"context"
+
+	"github.com/databricks/cli/libs/env"
+)
+
+// Product name constants
+const (
+	ClaudeCode = "claude-code"
+	GeminiCLI  = "gemini-cli"
+	Cursor     = "cursor"
+)
+
+// Environment variable constants
+const (
+	claudeCodeEnvVar  = "CLAUDECODE"
+	geminiCliEnvVar   = "GEMINI_CLI"
+	cursorAgentEnvVar = "CURSOR_AGENT"
+)
+
+// key is a package-local type for context keys
+type key int
+
+const (
+	productKey = key(1)
+)
+
+// detect performs the actual detection logic.
+// Returns product name string or empty string if detection is ambiguous.
+// Only returns a product if exactly one agent is detected.
+func detect(ctx context.Context) string {
+	var detected []string
+
+	if env.Get(ctx, claudeCodeEnvVar) != "" {
+		detected = append(detected, ClaudeCode)
+	}
+
+	if env.Get(ctx, geminiCliEnvVar) != "" {
+		detected = append(detected, GeminiCLI)
+	}
+
+	if env.Get(ctx, cursorAgentEnvVar) != "" {
+		detected = append(detected, Cursor)
+	}
+
+	// Only return a product if exactly one agent is detected
+	if len(detected) == 1 {
+		return detected[0]
+	}
+
+	return ""
+}
+
+// Detect detects the agent and stores it in context.
+// It returns a new context with the detection result set.
+func Detect(ctx context.Context) context.Context {
+	return context.WithValue(ctx, productKey, detect(ctx))
+}
+
+// Mock is a helper for tests to mock the detection result.
+func Mock(ctx context.Context, product string) context.Context {
+	return context.WithValue(ctx, productKey, product)
+}
+
+// Product returns the detected agent product name from context.
+// Returns empty string if no agent was detected.
+// Panics if called before Detect() or Mock().
+func Product(ctx context.Context) string {
+	v := ctx.Value(productKey)
+	if v == nil {
+		panic("agent.Product called without calling agent.Detect first")
+	}
+	return v.(string)
+}

libs/agent/agent_test.go

@@ -0,0 +1,104 @@
+package agent
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/cli/libs/env"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestDetect(t *testing.T) {
+	ctx := context.Background()
+	// Clear other agent env vars to ensure clean test environment
+	ctx = env.Set(ctx, geminiCliEnvVar, "")
+	ctx = env.Set(ctx, cursorAgentEnvVar, "")
+	ctx = env.Set(ctx, claudeCodeEnvVar, "1")
+
+	ctx = Detect(ctx)
+
+	assert.Equal(t, ClaudeCode, Product(ctx))
+}
+
+func TestProductCalledBeforeDetect(t *testing.T) {
+	ctx := context.Background()
+
+	require.Panics(t, func() {
+		Product(ctx)
+	})
+}
+
+func TestMock(t *testing.T) {
+	ctx := context.Background()
+	ctx = Mock(ctx, "test-agent")
+
+	assert.Equal(t, "test-agent", Product(ctx))
+}
+
+func TestDetectNoAgent(t *testing.T) {
+	ctx := context.Background()
+	ctx = env.Set(ctx, claudeCodeEnvVar, "")
+	ctx = env.Set(ctx, geminiCliEnvVar, "")
+	ctx = env.Set(ctx, cursorAgentEnvVar, "")
+
+	ctx = Detect(ctx)
+
+	assert.Equal(t, "", Product(ctx))
+}
+
+func TestDetectClaudeCode(t *testing.T) {
+	ctx := context.Background()
+	// Clear other agent env vars to ensure clean test environment
+	ctx = env.Set(ctx, geminiCliEnvVar, "")
+	ctx = env.Set(ctx, cursorAgentEnvVar, "")
+	ctx = env.Set(ctx, claudeCodeEnvVar, "1")
+
+	result := detect(ctx)
+	assert.Equal(t, ClaudeCode, result)
+}
+
+func TestDetectGeminiCLI(t *testing.T) {
+	ctx := context.Background()
+	// Clear other agent env vars to ensure clean test environment
+	ctx = env.Set(ctx, claudeCodeEnvVar, "")
+	ctx = env.Set(ctx, cursorAgentEnvVar, "")
+	ctx = env.Set(ctx, geminiCliEnvVar, "1")
+
+	result := detect(ctx)
+	assert.Equal(t, GeminiCLI, result)
+}
+
+func TestDetectCursor(t *testing.T) {
+	ctx := context.Background()
+	// Clear other agent env vars to ensure clean test environment
+	ctx = env.Set(ctx, claudeCodeEnvVar, "")
+	ctx = env.Set(ctx, geminiCliEnvVar, "")
+	ctx = env.Set(ctx, cursorAgentEnvVar, "1")
+
+	result := detect(ctx)
+	assert.Equal(t, Cursor, result)
+}
+
+func TestDetectMultipleAgents(t *testing.T) {
+	ctx := context.Background()
+	// Clear all agent env vars first
+	ctx = env.Set(ctx, cursorAgentEnvVar, "")
+	// If multiple agents are detected, return empty string
+	ctx = env.Set(ctx, claudeCodeEnvVar, "1")
+	ctx = env.Set(ctx, geminiCliEnvVar, "1")
+
+	result := detect(ctx)
+	assert.Equal(t, "", result)
+}
+
+func TestDetectMultipleAgentsAllThree(t *testing.T) {
+	ctx := context.Background()
+	// If all three agents are detected, return empty string
+	ctx = env.Set(ctx, claudeCodeEnvVar, "1")
+	ctx = env.Set(ctx, geminiCliEnvVar, "1")
+	ctx = env.Set(ctx, cursorAgentEnvVar, "1")
+
+	result := detect(ctx)
+	assert.Equal(t, "", result)
+}