diff --git a/acceptance/bundle/user_agent/test.toml b/acceptance/bundle/user_agent/test.toml index 9295fedc55..24f4290b2e 100644 --- a/acceptance/bundle/user_agent/test.toml +++ b/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 = '' diff --git a/cmd/root/root.go b/cmd/root/root.go index 39c36e86d3..0466bc4641 100644 --- a/cmd/root/root.go +++ b/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) diff --git a/cmd/root/user_agent_agent.go b/cmd/root/user_agent_agent.go new file mode 100644 index 0000000000..739b05bdee --- /dev/null +++ b/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) +} diff --git a/cmd/root/user_agent_agent_test.go b/cmd/root/user_agent_agent_test.go new file mode 100644 index 0000000000..e0021338dc --- /dev/null +++ b/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/") +} diff --git a/libs/agent/agent.go b/libs/agent/agent.go new file mode 100644 index 0000000000..36ff22ccdf --- /dev/null +++ b/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) +} diff --git a/libs/agent/agent_test.go b/libs/agent/agent_test.go new file mode 100644 index 0000000000..549f5eae60 --- /dev/null +++ b/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) +}