diff --git a/mapbox_tilesets/agent_detect.py b/mapbox_tilesets/agent_detect.py new file mode 100644 index 0000000..7f1b634 --- /dev/null +++ b/mapbox_tilesets/agent_detect.py @@ -0,0 +1,67 @@ +"""Detect the AI coding agent (if any) driving this CLI invocation.""" + +import os + +# (agent_id, [env_var, ...]) — table order is precedence order; the first entry with any +# of its env vars present wins. Presence is the only thing ever tested - a var's value is +# never read or compared against anything, for any entry. Even a var explicitly set to "" +# or whitespace counts as present. +# +# Ported from mapbox-sdk-js's `lib/helpers/agent-detect.js` (this repo's sibling +# implementation of the same allowlist - keep the two in sync). Canonical origin: +# HuggingFace's public `agent-harnesses.ts` registry. +# +# "vtcode" and "warp" used to require a specific value (VTCODE == "1", TERM_PROGRAM == +# "WarpTerminal") rather than mere presence. Since values are never checked, "warp" was +# dropped entirely: TERM_PROGRAM is set by most terminal emulators (iTerm2, Apple +# Terminal, VS Code, Hyper, ...), not just Warp, so an existence check on it would +# misidentify most terminal sessions as "warp". VTCODE has no such collision risk and +# stays as a plain presence check. +# +# The final entry, "custom-agent", is a catch-all for AI_AGENT/AGENT: these exist so an +# agent not on this list can still flag its presence, but we only ever check for them, +# never read their value - an arbitrary, unvalidated string must never be forwarded into +# telemetry as an "agent id". +_ALLOWLIST = [ + ("antigravity", ["ANTIGRAVITY_AGENT"]), + ("augment-cli", ["AUGMENT_AGENT"]), + ("cline", ["CLINE_ACTIVE"]), + ("cowork", ["CLAUDE_CODE_IS_COWORK"]), + ("claude-code", ["CLAUDECODE", "CLAUDE_CODE"]), + ("codex", ["CODEX_SANDBOX", "CODEX_CI", "CODEX_THREAD_ID"]), + ("crush", ["CRUSH"]), + ("gemini-cli", ["GEMINI_CLI"]), + ("github-copilot", ["COPILOT_MODEL", "COPILOT_ALLOW_ALL", "COPILOT_GITHUB_TOKEN"]), + ("goose", ["GOOSE_TERMINAL"]), + ("hermes-agent", ["HERMES_SESSION_ID"]), + ("kilo-code", ["KILOCODE_FEATURE"]), + ("kiro", ["AGENT_CONTEXT_OUT"]), + ("openclaw", ["OPENCLAW_SHELL"]), + ("opencode", ["OPENCODE_CLIENT"]), + ("pi", ["PI_CODING_AGENT"]), + ("replit", ["REPL_ID"]), + ("trae", ["TRAE_AI_SHELL_ID"]), + ("vtcode", ["VTCODE"]), + ("zed", ["ZED_TERM"]), + ("cursor-cli", ["CURSOR_AGENT"]), + ("cursor", ["CURSOR_TRACE_ID"]), + ("custom-agent", ["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, env_vars in _ALLOWLIST: + for var in env_vars: + if var in os.environ: + return agent_id + + return None diff --git a/mapbox_tilesets/utils.py b/mapbox_tilesets/utils.py index 8f8e2fa..9da8a93 100644 --- a/mapbox_tilesets/utils.py +++ b/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 diff --git a/tests/test_agent_detect.py b/tests/test_agent_detect.py new file mode 100644 index 0000000..d2c628d --- /dev/null +++ b/tests/test_agent_detect.py @@ -0,0 +1,93 @@ +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_was_dropped_term_program_is_not_a_safe_existence_only_signal(): + # TERM_PROGRAM is set by most terminal emulators, not just Warp, so it's + # not on the allowlist at all now that presence is the only check. + assert _detect_with_env({"TERM_PROGRAM": "WarpTerminal"}) is None + assert _detect_with_env({"TERM_PROGRAM": "iTerm.app"}) is None + + +def test_vtcode_matches_on_presence_alone_regardless_of_value(): + assert _detect_with_env({"VTCODE": "1"}) == "vtcode" + assert _detect_with_env({"VTCODE": "0"}) == "vtcode" + assert _detect_with_env({"VTCODE": "true"}) == "vtcode" + + +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": "my-cool-tool"}) == "custom-agent" + + +def test_fallback_agent_used_when_no_harness_or_ai_agent_match(): + assert _detect_with_env({"AGENT": "my-cool-tool"}) == "custom-agent" + + +def test_fallback_ai_agent_takes_precedence_over_agent(): + # Same result either way - table order still decides precedence, but the + # value of AI_AGENT/AGENT is never read, only their presence. + assert _detect_with_env({"AI_AGENT": "first", "AGENT": "second"}) == "custom-agent" + + +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_only_requires_existence_blank_value_still_counts(): + # Existence is all that matters - a var set to "" or whitespace still + # counts as present, for both a harness var and the AI_AGENT/AGENT + # catch-all. + assert _detect_with_env({"CLAUDECODE": ""}) == "claude-code" + assert _detect_with_env({"CLAUDECODE": " "}) == "claude-code" + assert _detect_with_env({"AI_AGENT": ""}) == "custom-agent" + assert _detect_with_env({"AI_AGENT": " "}) == "custom-agent" + + +def test_fallback_value_is_never_forwarded_even_if_header_unsafe(): + # The fallback value must never reach the User-Agent header at all - it is + # never read, only checked for presence. + assert _detect_with_env({"AI_AGENT": "foo\nbar: injected"}) == "custom-agent" diff --git a/tests/test_utils.py b/tests/test_utils.py index f3d79c8..05a5f04 100644 --- a/tests/test_utils.py +++ b/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)