Skip to content

Comments

[PECOBLR-1928] Add AI coding agent detection to User-Agent header#326

Open
vikrantpuppala wants to merge 1 commit intodatabricks:mainfrom
vikrantpuppala:agent-detection
Open

[PECOBLR-1928] Add AI coding agent detection to User-Agent header#326
vikrantpuppala wants to merge 1 commit intodatabricks:mainfrom
vikrantpuppala:agent-detection

Conversation

@vikrantpuppala
Copy link
Collaborator

Summary

  • Adds internal/agent package 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 processes
  • Integrates detection into InitThriftClient to append agent/<product> to the User-Agent header
  • Uses exactly-one detection rule: if zero or multiple agent env vars are set, no agent is attributed (avoids ambiguity)

Approach

Mirrors the implementation in databricks/cli#4287 and aligns with the latest agent list in libs/agent/agent.go.

Agent Product String Environment Variable
Google Antigravity antigravity ANTIGRAVITY_AGENT
Claude Code claude-code CLAUDECODE
Cline cline CLINE_ACTIVE
OpenAI Codex codex CODEX_CI
Cursor cursor CURSOR_AGENT
Gemini CLI gemini-cli GEMINI_CLI
OpenCode opencode OPENCODE

Adding a new agent requires only a new constant and a new entry in knownAgents.

Changes

  • New: internal/agent/agent.go — environment-variable-based agent detection with injectable env lookup for testability
  • New: internal/agent/agent_test.go — 11 test cases covering all agents, no agent, multiple agents, empty values, and real os.Getenv
  • Modified: internal/client/client.go — calls agent.Detect() when building User-Agent in InitThriftClient

Test plan

  • internal/agent — 11 tests pass
  • internal/client — all existing tests continue to pass
  • Manual: verified User-Agent contains agent/claude-code when run from Claude Code via GODEBUG=http2debug=2
    http2: Transport encoding header "user-agent" = "godatabrickssqlconnector/1.10.0 agent/claude-code"
    
  • Executed SELECT 1 successfully against dogfood warehouse with the new header

🤖 Generated with Claude Code

Copilot AI review requested due to automatic review settings February 24, 2026 07:30
Copy link

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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/agent package 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.

@vikrantpuppala vikrantpuppala changed the title Add AI coding agent detection to User-Agent header [PECOBLR-1928] Add AI coding agent detection to User-Agent header Feb 24, 2026
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants