[PECOBLR-1928] Add AI coding agent detection to User-Agent header#326
Open
vikrantpuppala wants to merge 1 commit intodatabricks:mainfrom
Open
[PECOBLR-1928] Add AI coding agent detection to User-Agent header#326vikrantpuppala wants to merge 1 commit intodatabricks:mainfrom
vikrantpuppala wants to merge 1 commit intodatabricks:mainfrom
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds AI coding agent detection to the User-Agent header by checking for well-known environment variables set by 7 different AI coding agents. The detection uses an "exactly-one" rule to avoid ambiguous attribution when multiple agent environments overlap. The implementation mirrors the approach from databricks/cli#4287.
Changes:
- Added new
internal/agentpackage with environment-variable-based agent detection logic and comprehensive test coverage - Integrated agent detection into the User-Agent header construction in
InitThriftClient
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
| internal/agent/agent.go | New package implementing AI coding agent detection via environment variables with injectable env lookup for testability |
| internal/agent/agent_test.go | Comprehensive test suite with 11 test cases covering all agents, edge cases, and real os.Getenv usage |
| internal/client/client.go | Integration of agent detection into User-Agent header construction, appending agent/<product> when detected |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Detect when the Go SQL 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> Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
c0f7406 to
1ededd0
Compare
simonfaltum
approved these changes
Feb 24, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
internal/agentpackage that detects 7 AI coding agents (Claude Code, Cursor, Gemini CLI, Cline, Codex, OpenCode, Antigravity) by checking well-known environment variables they set in spawned shell processesInitThriftClientto appendagent/<product>to the User-Agent headerApproach
Mirrors the implementation in databricks/cli#4287 and aligns with the latest agent list in
libs/agent/agent.go.antigravityANTIGRAVITY_AGENTclaude-codeCLAUDECODEclineCLINE_ACTIVEcodexCODEX_CIcursorCURSOR_AGENTgemini-cliGEMINI_CLIopencodeOPENCODEAdding a new agent requires only a new constant and a new entry in
knownAgents.Changes
internal/agent/agent.go— environment-variable-based agent detection with injectable env lookup for testabilityinternal/agent/agent_test.go— 11 test cases covering all agents, no agent, multiple agents, empty values, and realos.Getenvinternal/client/client.go— callsagent.Detect()when building User-Agent inInitThriftClientTest plan
internal/agent— 11 tests passinternal/client— all existing tests continue to passagent/claude-codewhen run from Claude Code viaGODEBUG=http2debug=2SELECT 1successfully against dogfood warehouse with the new header🤖 Generated with Claude Code