From 96fac3ca1559a26a6f65b8ee610a30ed2517c16b Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Tue, 24 Feb 2026 06:40:32 +0000 Subject: [PATCH] Add AI coding agent detection to User-Agent header Detect when the JDBC driver is invoked by an AI coding agent (e.g. Claude Code, Cursor, Gemini CLI) by checking well-known environment variables, and append `agent/` to the User-Agent string. This enables Databricks to understand how much driver usage originates from AI coding agents. Detection only succeeds when exactly one agent is detected to avoid ambiguous attribution. Mirrors the approach in databricks/cli#4287. Co-Authored-By: Claude Opus 4.6 Signed-off-by: Vikrant Puppala --- NEXT_CHANGELOG.md | 1 + .../jdbc/common/util/AgentDetector.java | 78 ++++++++++++++++++ .../jdbc/common/util/UserAgentManager.java | 13 +++ .../jdbc/common/util/AgentDetectorTest.java | 79 +++++++++++++++++++ 4 files changed, 171 insertions(+) create mode 100644 src/main/java/com/databricks/jdbc/common/util/AgentDetector.java create mode 100644 src/test/java/com/databricks/jdbc/common/util/AgentDetectorTest.java diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 499d04dca..05fecadde 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -3,6 +3,7 @@ ## [Unreleased] ### Added +- Added AI coding agent detection to the User-Agent header. When the driver is invoked by a known AI coding agent (e.g. Claude Code, Cursor, Gemini CLI), `agent/` is appended to the User-Agent string. - Added streaming prefetch mode for Thrift inline results (columnar and Arrow) with background batch prefetching and configurable sliding window for improved throughput. - Added `EnableInlineStreaming` connection parameter to enable/disable streaming mode (default: enabled). - Added `ThriftMaxBatchesInMemory` connection parameter to control the sliding window size for streaming (default: 3). diff --git a/src/main/java/com/databricks/jdbc/common/util/AgentDetector.java b/src/main/java/com/databricks/jdbc/common/util/AgentDetector.java new file mode 100644 index 000000000..e78a8205c --- /dev/null +++ b/src/main/java/com/databricks/jdbc/common/util/AgentDetector.java @@ -0,0 +1,78 @@ +package com.databricks.jdbc.common.util; + +import java.util.ArrayList; +import java.util.List; +import java.util.function.Function; + +/** + * Detects whether the JDBC driver is being invoked by an AI coding agent by checking for well-known + * environment variables that agents set in their spawned shell processes. + * + *

Detection only succeeds when exactly one agent environment variable is present, to avoid + * ambiguous attribution when multiple agent environments overlap. + * + *

Adding a new agent requires only a new constant and a new entry in {@link #KNOWN_AGENTS}. + * + *

References for each environment variable: + * + *

    + *
  • ANTIGRAVITY_AGENT: Closed source. Google Antigravity sets this variable. + *
  • CLAUDECODE: https://github.com/anthropics/claude-code (sets CLAUDECODE=1) + *
  • CLINE_ACTIVE: https://github.com/cline/cline (shipped in v3.24.0) + *
  • CODEX_CI: https://github.com/openai/codex (part of UNIFIED_EXEC_ENV array in codex-rs) + *
  • CURSOR_AGENT: Closed source. Referenced in a gist by johnlindquist. + *
  • GEMINI_CLI: https://google-gemini.github.io/gemini-cli/docs/tools/shell.html (sets + * GEMINI_CLI=1) + *
  • OPENCODE: https://github.com/opencode-ai/opencode (sets OPENCODE=1) + *
+ */ +public class AgentDetector { + + public static final String ANTIGRAVITY = "antigravity"; + public static final String CLAUDE_CODE = "claude-code"; + public static final String CLINE = "cline"; + public static final String CODEX = "codex"; + public static final String CURSOR = "cursor"; + public static final String GEMINI_CLI = "gemini-cli"; + public static final String OPEN_CODE = "opencode"; + + static final String[][] KNOWN_AGENTS = { + {"ANTIGRAVITY_AGENT", ANTIGRAVITY}, + {"CLAUDECODE", CLAUDE_CODE}, + {"CLINE_ACTIVE", CLINE}, + {"CODEX_CI", CODEX}, + {"CURSOR_AGENT", CURSOR}, + {"GEMINI_CLI", GEMINI_CLI}, + {"OPENCODE", OPEN_CODE}, + }; + + /** + * Detects which AI coding agent (if any) is driving the current process. + * + * @return the agent product string if exactly one agent is detected, or an empty string otherwise + */ + public static String detect() { + return detect(System::getenv); + } + + /** + * Detects which AI coding agent (if any) is present, using the provided function to look up + * environment variables. This overload exists for testability. + * + * @param envLookup function that returns the value of an environment variable, or null if unset + * @return the agent product string if exactly one agent is detected, or an empty string otherwise + */ + static String detect(Function envLookup) { + List detected = new ArrayList<>(); + for (String[] entry : KNOWN_AGENTS) { + String value = envLookup.apply(entry[0]); + if (value != null && !value.isEmpty()) { + detected.add(entry[1]); + } + } + if (detected.size() == 1) { + return detected.get(0); + } + return ""; + } +} diff --git a/src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java b/src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java index 14ac19f05..e6593c276 100644 --- a/src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java +++ b/src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java @@ -16,6 +16,7 @@ public class UserAgentManager { public static final String USER_AGENT_SEA_CLIENT = "SQLExecHttpClient"; public static final String USER_AGENT_THRIFT_CLIENT = "THttpClient"; private static final String VERSION_FILLER = "version"; + private static final String AGENT_KEY = "agent"; /** * Parse custom user agent string into name and version components. @@ -62,6 +63,12 @@ public static void setUserAgent(IDatabricksConnectionContext connectionContext) } } } + + // Detect AI coding agent and append to user agent + String agentProduct = AgentDetector.detect(); + if (!agentProduct.isEmpty()) { + UserAgent.withOtherInfo(AGENT_KEY, agentProduct); + } } /** @@ -106,6 +113,12 @@ public static String buildUserAgentForConnectorService( } } + // Detect AI coding agent and append to user agent + String agentProduct = AgentDetector.detect(); + if (!agentProduct.isEmpty()) { + userAgent.append(" ").append(AGENT_KEY).append("/").append(agentProduct); + } + return userAgent.toString(); } diff --git a/src/test/java/com/databricks/jdbc/common/util/AgentDetectorTest.java b/src/test/java/com/databricks/jdbc/common/util/AgentDetectorTest.java new file mode 100644 index 000000000..f576ef3a3 --- /dev/null +++ b/src/test/java/com/databricks/jdbc/common/util/AgentDetectorTest.java @@ -0,0 +1,79 @@ +package com.databricks.jdbc.common.util; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import java.util.HashMap; +import java.util.Map; +import java.util.stream.Stream; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +public class AgentDetectorTest { + + /** Creates an env lookup function that returns values from the given map. */ + private static java.util.function.Function envWith(Map env) { + return env::get; + } + + @ParameterizedTest + @MethodSource("singleAgentCases") + void testDetectsSingleAgent(String envVar, String expectedProduct) { + Map env = new HashMap<>(); + env.put(envVar, "1"); + assertEquals(expectedProduct, AgentDetector.detect(envWith(env))); + } + + static Stream singleAgentCases() { + return Stream.of( + Arguments.of("ANTIGRAVITY_AGENT", AgentDetector.ANTIGRAVITY), + Arguments.of("CLAUDECODE", AgentDetector.CLAUDE_CODE), + Arguments.of("CLINE_ACTIVE", AgentDetector.CLINE), + Arguments.of("CODEX_CI", AgentDetector.CODEX), + Arguments.of("CURSOR_AGENT", AgentDetector.CURSOR), + Arguments.of("GEMINI_CLI", AgentDetector.GEMINI_CLI), + Arguments.of("OPENCODE", AgentDetector.OPEN_CODE)); + } + + @Test + void testReturnsEmptyWhenNoAgentDetected() { + Map env = new HashMap<>(); + assertEquals("", AgentDetector.detect(envWith(env))); + } + + @Test + void testReturnsEmptyWhenMultipleAgentsDetected() { + Map env = new HashMap<>(); + env.put("CLAUDECODE", "1"); + env.put("CURSOR_AGENT", "1"); + assertEquals("", AgentDetector.detect(envWith(env))); + } + + @Test + void testIgnoresEmptyEnvVarValues() { + Map env = new HashMap<>(); + env.put("CLAUDECODE", ""); + assertEquals("", AgentDetector.detect(envWith(env))); + } + + @Test + void testIgnoresNullEnvVarValues() { + Map env = new HashMap<>(); + env.put("CLAUDECODE", null); + assertEquals("", AgentDetector.detect(envWith(env))); + } + + @Test + void testAllKnownAgentsAreCovered() { + // Verify every entry in KNOWN_AGENTS can be detected individually + for (String[] entry : AgentDetector.KNOWN_AGENTS) { + Map env = new HashMap<>(); + env.put(entry[0], "1"); + assertEquals( + entry[1], + AgentDetector.detect(envWith(env)), + "Agent with env var " + entry[0] + " should be detected as " + entry[1]); + } + } +}