Skip to content

Commit fd74f70

Browse files
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/<product>` 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 <noreply@anthropic.com>
1 parent 8e23bc4 commit fd74f70

File tree

3 files changed

+170
-0
lines changed

3 files changed

+170
-0
lines changed
Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
package com.databricks.jdbc.common.util;
2+
3+
import java.util.ArrayList;
4+
import java.util.List;
5+
import java.util.function.Function;
6+
7+
/**
8+
* Detects whether the JDBC driver is being invoked by an AI coding agent by checking for well-known
9+
* environment variables that agents set in their spawned shell processes.
10+
*
11+
* <p>Detection only succeeds when exactly one agent environment variable is present, to avoid
12+
* ambiguous attribution when multiple agent environments overlap.
13+
*
14+
* <p>Adding a new agent requires only a new constant and a new entry in {@link #KNOWN_AGENTS}.
15+
*
16+
* <p>References for each environment variable:
17+
*
18+
* <ul>
19+
* <li>ANTIGRAVITY_AGENT: Closed source. Google Antigravity sets this variable.
20+
* <li>CLAUDECODE: https://github.com/anthropics/claude-code (sets CLAUDECODE=1)
21+
* <li>CLINE_ACTIVE: https://github.com/cline/cline (shipped in v3.24.0)
22+
* <li>CODEX_CI: https://github.com/openai/codex (part of UNIFIED_EXEC_ENV array in codex-rs)
23+
* <li>CURSOR_AGENT: Closed source. Referenced in a gist by johnlindquist.
24+
* <li>GEMINI_CLI: https://google-gemini.github.io/gemini-cli/docs/tools/shell.html (sets
25+
* GEMINI_CLI=1)
26+
* <li>OPENCODE: https://github.com/opencode-ai/opencode (sets OPENCODE=1)
27+
* </ul>
28+
*/
29+
public class AgentDetector {
30+
31+
public static final String ANTIGRAVITY = "antigravity";
32+
public static final String CLAUDE_CODE = "claude-code";
33+
public static final String CLINE = "cline";
34+
public static final String CODEX = "codex";
35+
public static final String CURSOR = "cursor";
36+
public static final String GEMINI_CLI = "gemini-cli";
37+
public static final String OPEN_CODE = "opencode";
38+
39+
static final String[][] KNOWN_AGENTS = {
40+
{"ANTIGRAVITY_AGENT", ANTIGRAVITY},
41+
{"CLAUDECODE", CLAUDE_CODE},
42+
{"CLINE_ACTIVE", CLINE},
43+
{"CODEX_CI", CODEX},
44+
{"CURSOR_AGENT", CURSOR},
45+
{"GEMINI_CLI", GEMINI_CLI},
46+
{"OPENCODE", OPEN_CODE},
47+
};
48+
49+
/**
50+
* Detects which AI coding agent (if any) is driving the current process.
51+
*
52+
* @return the agent product string if exactly one agent is detected, or an empty string otherwise
53+
*/
54+
public static String detect() {
55+
return detect(System::getenv);
56+
}
57+
58+
/**
59+
* Detects which AI coding agent (if any) is present, using the provided function to look up
60+
* environment variables. This overload exists for testability.
61+
*
62+
* @param envLookup function that returns the value of an environment variable, or null if unset
63+
* @return the agent product string if exactly one agent is detected, or an empty string otherwise
64+
*/
65+
static String detect(Function<String, String> envLookup) {
66+
List<String> detected = new ArrayList<>();
67+
for (String[] entry : KNOWN_AGENTS) {
68+
String value = envLookup.apply(entry[0]);
69+
if (value != null && !value.isEmpty()) {
70+
detected.add(entry[1]);
71+
}
72+
}
73+
if (detected.size() == 1) {
74+
return detected.get(0);
75+
}
76+
return "";
77+
}
78+
}

src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ public class UserAgentManager {
1616
public static final String USER_AGENT_SEA_CLIENT = "SQLExecHttpClient";
1717
public static final String USER_AGENT_THRIFT_CLIENT = "THttpClient";
1818
private static final String VERSION_FILLER = "version";
19+
private static final String AGENT_KEY = "agent";
1920

2021
/**
2122
* Parse custom user agent string into name and version components.
@@ -62,6 +63,12 @@ public static void setUserAgent(IDatabricksConnectionContext connectionContext)
6263
}
6364
}
6465
}
66+
67+
// Detect AI coding agent and append to user agent
68+
String agentProduct = AgentDetector.detect();
69+
if (!agentProduct.isEmpty()) {
70+
UserAgent.withOtherInfo(AGENT_KEY, agentProduct);
71+
}
6572
}
6673

6774
/**
@@ -106,6 +113,12 @@ public static String buildUserAgentForConnectorService(
106113
}
107114
}
108115

116+
// Detect AI coding agent and append to user agent
117+
String agentProduct = AgentDetector.detect();
118+
if (!agentProduct.isEmpty()) {
119+
userAgent.append(" ").append(AGENT_KEY).append("/").append(agentProduct);
120+
}
121+
109122
return userAgent.toString();
110123
}
111124

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
package com.databricks.jdbc.common.util;
2+
3+
import static org.junit.jupiter.api.Assertions.assertEquals;
4+
5+
import java.util.HashMap;
6+
import java.util.Map;
7+
import java.util.stream.Stream;
8+
import org.junit.jupiter.api.Test;
9+
import org.junit.jupiter.params.ParameterizedTest;
10+
import org.junit.jupiter.params.provider.Arguments;
11+
import org.junit.jupiter.params.provider.MethodSource;
12+
13+
public class AgentDetectorTest {
14+
15+
/** Creates an env lookup function that returns values from the given map. */
16+
private static java.util.function.Function<String, String> envWith(Map<String, String> env) {
17+
return env::get;
18+
}
19+
20+
@ParameterizedTest
21+
@MethodSource("singleAgentCases")
22+
void testDetectsSingleAgent(String envVar, String expectedProduct) {
23+
Map<String, String> env = new HashMap<>();
24+
env.put(envVar, "1");
25+
assertEquals(expectedProduct, AgentDetector.detect(envWith(env)));
26+
}
27+
28+
static Stream<Arguments> singleAgentCases() {
29+
return Stream.of(
30+
Arguments.of("ANTIGRAVITY_AGENT", AgentDetector.ANTIGRAVITY),
31+
Arguments.of("CLAUDECODE", AgentDetector.CLAUDE_CODE),
32+
Arguments.of("CLINE_ACTIVE", AgentDetector.CLINE),
33+
Arguments.of("CODEX_CI", AgentDetector.CODEX),
34+
Arguments.of("CURSOR_AGENT", AgentDetector.CURSOR),
35+
Arguments.of("GEMINI_CLI", AgentDetector.GEMINI_CLI),
36+
Arguments.of("OPENCODE", AgentDetector.OPEN_CODE));
37+
}
38+
39+
@Test
40+
void testReturnsEmptyWhenNoAgentDetected() {
41+
Map<String, String> env = new HashMap<>();
42+
assertEquals("", AgentDetector.detect(envWith(env)));
43+
}
44+
45+
@Test
46+
void testReturnsEmptyWhenMultipleAgentsDetected() {
47+
Map<String, String> env = new HashMap<>();
48+
env.put("CLAUDECODE", "1");
49+
env.put("CURSOR_AGENT", "1");
50+
assertEquals("", AgentDetector.detect(envWith(env)));
51+
}
52+
53+
@Test
54+
void testIgnoresEmptyEnvVarValues() {
55+
Map<String, String> env = new HashMap<>();
56+
env.put("CLAUDECODE", "");
57+
assertEquals("", AgentDetector.detect(envWith(env)));
58+
}
59+
60+
@Test
61+
void testIgnoresNullEnvVarValues() {
62+
Map<String, String> env = new HashMap<>();
63+
env.put("CLAUDECODE", null);
64+
assertEquals("", AgentDetector.detect(envWith(env)));
65+
}
66+
67+
@Test
68+
void testAllKnownAgentsAreCovered() {
69+
// Verify every entry in KNOWN_AGENTS can be detected individually
70+
for (String[] entry : AgentDetector.KNOWN_AGENTS) {
71+
Map<String, String> env = new HashMap<>();
72+
env.put(entry[0], "1");
73+
assertEquals(
74+
entry[1],
75+
AgentDetector.detect(envWith(env)),
76+
"Agent with env var " + entry[0] + " should be detected as " + entry[1]);
77+
}
78+
}
79+
}

0 commit comments

Comments
 (0)