Append agent/<id> to the User-Agent when driven by an AI coding agent

mapbox_tilesets/agent_detect.py

@@ -0,0 +1,75 @@
+"""Detect the AI coding agent (if any) driving this CLI invocation."""
+
+import os
+import re
+
+# A safe charset for an agent id placed into an HTTP header: env vars are not
+# validated by whoever sets them, so a value like "foo\nbar: injected" must be
+# rejected here rather than surfacing as an unhandled requests.InvalidHeader.
+_SAFE_FALLBACK_ID = re.compile(r"^[\w.\-]{1,64}$")
+
+# (agent_id, [(env_var, expected_value_or_None), ...]) — table order is precedence order;
+# the first entry with any matching condition wins. expected_value None => presence check
+# (key exists in os.environ with a non-empty, non-whitespace value); otherwise an
+# exact-equality check.
+_ALLOWLIST = [
+    ("antigravity", [("ANTIGRAVITY_AGENT", None)]),
+    ("augment-cli", [("AUGMENT_AGENT", None)]),
+    ("cline", [("CLINE_ACTIVE", None)]),
+    ("cowork", [("CLAUDE_CODE_IS_COWORK", None)]),
+    ("claude-code", [("CLAUDECODE", None), ("CLAUDE_CODE", None)]),
+    ("codex", [("CODEX_SANDBOX", None), ("CODEX_CI", None), ("CODEX_THREAD_ID", None)]),
+    ("crush", [("CRUSH", None)]),
+    ("gemini-cli", [("GEMINI_CLI", None)]),
+    (
+        "github-copilot",
+        [
+            ("COPILOT_MODEL", None),
+            ("COPILOT_ALLOW_ALL", None),
+            ("COPILOT_GITHUB_TOKEN", None),
+        ],
+    ),
+    ("goose", [("GOOSE_TERMINAL", None)]),
+    ("hermes-agent", [("HERMES_SESSION_ID", None)]),
+    ("kilo-code", [("KILOCODE_FEATURE", None)]),
+    ("kiro", [("AGENT_CONTEXT_OUT", None)]),
+    ("openclaw", [("OPENCLAW_SHELL", None)]),
+    ("opencode", [("OPENCODE_CLIENT", None)]),
+    ("pi", [("PI_CODING_AGENT", None)]),
+    ("replit", [("REPL_ID", None)]),
+    ("trae", [("TRAE_AI_SHELL_ID", None)]),
+    ("vtcode", [("VTCODE", "1")]),
+    ("warp", [("TERM_PROGRAM", "WarpTerminal")]),
+    ("zed", [("ZED_TERM", None)]),
+    ("cursor-cli", [("CURSOR_AGENT", None)]),
+    ("cursor", [("CURSOR_TRACE_ID", None)]),
+]
+
+# Checked only if nothing in _ALLOWLIST matched. First one with a non-empty value wins.
+_FALLBACK_VARS = ("AI_AGENT", "AGENT")
+
+
+def detect_agent():
+    """Detect the AI coding agent driving this CLI from environment variables.
+
+    Never reads or logs the full environment - only the matched id is used.
+
+    Returns
+    -------
+    str or None
+        The detected agent id, or None when no agent indicator is present.
+    """
+    for agent_id, conditions in _ALLOWLIST:
+        for var, expected in conditions:
+            if expected is None:
+                if os.environ.get(var, "").strip():
+                    return agent_id
+            elif os.environ.get(var) == expected:
+                return agent_id
+
+    for var in _FALLBACK_VARS:
+        value = os.environ.get(var, "").strip()
+        if value and _SAFE_FALLBACK_ID.match(value):
+            return value
+
+    return None

mapbox_tilesets/utils.py

@@ -12,6 +12,8 @@
 import geojson
 import json
 
+from mapbox_tilesets.agent_detect import detect_agent
+
 
 def load_module(modulename):
     """Dynamically imports a module and throws a readable exception if not found"""
@@ -51,7 +53,11 @@ def _get_session(
 ):
     """Get a configured session"""
     s = Session()
-    s.headers.update({"user-agent": "{}/{}".format(application, version)})
+    user_agent = "{}/{}".format(application, version)
+    agent = detect_agent()
+    if agent:
+        user_agent = "{} agent/{}".format(user_agent, agent)
+    s.headers.update({"user-agent": user_agent})
     return s
 
 

tests/test_agent_detect.py

@@ -0,0 +1,104 @@
+import os
+from unittest import mock
+
+from mapbox_tilesets.agent_detect import detect_agent
+
+
+def _detect_with_env(env):
+    # clear=True: these tests run inside various AI-agent shells (e.g. Claude
+    # Code sets CLAUDECODE), so the ambient environment must not leak in.
+    with mock.patch.dict(os.environ, env, clear=True):
+        return detect_agent()
+
+
+def test_no_indicators_returns_none():
+    assert _detect_with_env({}) is None
+
+
+def test_harness_var_wins_over_ai_agent_fallback():
+    # Harness-specific vars take precedence over AI_AGENT/AGENT even when both
+    # are present at once.
+    assert (
+        _detect_with_env({"CLAUDECODE": "1", "AI_AGENT": "something-else"})
+        == "claude-code"
+    )
+
+
+def test_codex_and_claude_code_are_distinct():
+    assert _detect_with_env({"CODEX_THREAD_ID": "abc"}) == "codex"
+    assert _detect_with_env({"CLAUDECODE": "1"}) == "claude-code"
+    assert _detect_with_env({"CLAUDE_CODE": "1"}) == "claude-code"
+
+
+def test_codex_matches_on_any_of_its_vars():
+    assert _detect_with_env({"CODEX_SANDBOX": "1"}) == "codex"
+    assert _detect_with_env({"CODEX_CI": "1"}) == "codex"
+
+
+def test_warp_requires_exact_value_match():
+    assert _detect_with_env({"TERM_PROGRAM": "WarpTerminal"}) == "warp"
+    assert _detect_with_env({"TERM_PROGRAM": "iTerm.app"}) is None
+
+
+def test_vtcode_requires_exact_value_match():
+    assert _detect_with_env({"VTCODE": "1"}) == "vtcode"
+    assert _detect_with_env({"VTCODE": "0"}) is None
+    assert _detect_with_env({"VTCODE": "true"}) is None
+
+
+def test_table_order_precedence_among_harness_vars():
+    # antigravity is earlier in the table than cursor - it should win when
+    # both indicators are present.
+    assert (
+        _detect_with_env({"CURSOR_AGENT": "1", "ANTIGRAVITY_AGENT": "1"})
+        == "antigravity"
+    )
+
+
+def test_fallback_ai_agent_used_when_no_harness_match():
+    assert _detect_with_env({"AI_AGENT": "custom-agent"}) == "custom-agent"
+
+
+def test_fallback_agent_used_when_no_harness_or_ai_agent_match():
+    assert _detect_with_env({"AGENT": "custom-agent"}) == "custom-agent"
+
+
+def test_fallback_ai_agent_takes_precedence_over_agent():
+    assert _detect_with_env({"AI_AGENT": "first", "AGENT": "second"}) == "first"
+
+
+def test_fallback_empty_or_whitespace_value_returns_none():
+    assert _detect_with_env({"AI_AGENT": ""}) is None
+    assert _detect_with_env({"AI_AGENT": "   "}) is None
+    assert _detect_with_env({"AI_AGENT": "", "AGENT": "still-empty-check"}) == (
+        "still-empty-check"
+    )
+
+
+def test_github_copilot_matches_on_any_of_its_vars():
+    assert _detect_with_env({"COPILOT_MODEL": "gpt"}) == "github-copilot"
+    assert _detect_with_env({"COPILOT_ALLOW_ALL": "1"}) == "github-copilot"
+    assert _detect_with_env({"COPILOT_GITHUB_TOKEN": "abc"}) == "github-copilot"
+
+
+def test_presence_check_requires_non_empty_value():
+    # A harness var set to "" or whitespace is treated the same as unset,
+    # matching the fallback's empty-value handling.
+    assert _detect_with_env({"CLAUDECODE": ""}) is None
+    assert _detect_with_env({"CLAUDECODE": "   "}) is None
+
+
+def test_fallback_rejects_header_unsafe_characters():
+    # A fallback value must never reach the User-Agent header unsanitized -
+    # a newline would otherwise crash every request with InvalidHeader.
+    assert _detect_with_env({"AI_AGENT": "foo\nbar: injected"}) is None
+    assert _detect_with_env({"AI_AGENT": "has spaces"}) is None
+
+
+def test_fallback_rejects_unsafe_value_but_falls_through_to_next_var():
+    assert _detect_with_env({"AI_AGENT": "foo\nbar", "AGENT": "safe-id"}) == "safe-id"
+
+
+def test_fallback_rejects_overlong_value():
+    assert _detect_with_env({"AI_AGENT": "a" * 65}) is None
+    assert _detect_with_env({"AI_AGENT": "a" * 64}) == "a" * 64

tests/test_utils.py

@@ -1,6 +1,7 @@
 import os
 import pytest
 import json
+from unittest import mock
 from mapbox_tilesets.utils import (
     _get_api,
     _get_session,
@@ -22,11 +23,20 @@ def test_get_api():
 
 
 def test_get_session():
-    s = _get_session("my_application", "1.0.0")
+    # Cleared so an ambient agent indicator (e.g. this test running inside
+    # Claude Code, where CLAUDECODE is set) can't leak into the assertion.
+    with mock.patch.dict(os.environ, {}, clear=True):
+        s = _get_session("my_application", "1.0.0")
     assert "user-agent" in s.headers
     assert s.headers["user-agent"] == "my_application/1.0.0"
 
 
+def test_get_session_appends_detected_agent():
+    with mock.patch.dict(os.environ, {"CLAUDECODE": "1"}, clear=True):
+        s = _get_session("my_application", "1.0.0")
+    assert s.headers["user-agent"] == "my_application/1.0.0 agent/claude-code"
+
+
 def test_get_token_parameter():
     token = "token-parameter"
     assert token == _get_token(token)