From c299cef356edefc34a8417dd14d9db5835cab2f1 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Thu, 15 Jan 2026 11:15:23 +0100 Subject: [PATCH] Add agent detection library MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds libs/agent package to detect agents (Claude Code, Gemini CLI, Cursor) via environment variables. Used in user agent tracking and available for telemetry and conditional behavior. Detection only succeeds when exactly one agent environment variable is set, preventing ambiguous detection in environments with multiple agents. Implementation: - libs/agent: Detection logic with Detect(), Product(), and Mock() - cmd/root: Integrated agent detection into command execution - Test config: Clear agent env vars in acceptance tests to prevent interference with test outputs Environment variable sources: - CLAUDECODE=1: https://github.com/anthropics/claude-code/issues/531 - GEMINI_CLI=1: https://github.com/google-gemini/gemini-cli/blob/main/docs/cli/commands.md - CURSOR_AGENT=1: https://forum.cursor.com/t/cursor-cli-is-not-setting-cursor-agent-1-environment-variable-while-executing-bash-commands/132427 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Sonnet 4.5 --- acceptance/bundle/user_agent/test.toml | 4 + cmd/root/root.go | 5 ++ cmd/root/user_agent_agent.go | 27 +++++++ cmd/root/user_agent_agent_test.go | 42 ++++++++++ libs/agent/agent.go | 76 ++++++++++++++++++ libs/agent/agent_test.go | 104 +++++++++++++++++++++++++ 6 files changed, 258 insertions(+) create mode 100644 cmd/root/user_agent_agent.go create mode 100644 cmd/root/user_agent_agent_test.go create mode 100644 libs/agent/agent.go create mode 100644 libs/agent/agent_test.go 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) +}