From da281ceb99b43866ff2793f78819925199421208 Mon Sep 17 00:00:00 2001 From: ctufts Date: Fri, 17 Jul 2026 16:50:32 -0400 Subject: [PATCH 1/2] Append agent/ to User-Agent when driven by an AI coding agent Detects the calling AI coding agent from a hardcoded env-var allowlist and appends " agent/" to the tilesets Session User-Agent when one is found. No server dependency: the header is already logged by CloudFront, and the agent's environment is directly visible to this CLI's subprocess. --- mapbox_tilesets/agent_detect.py | 68 +++++++++++++++++++++++++++ mapbox_tilesets/utils.py | 8 +++- tests/test_agent_detect.py | 81 +++++++++++++++++++++++++++++++++ tests/test_utils.py | 12 ++++- 4 files changed, 167 insertions(+), 2 deletions(-) create mode 100644 mapbox_tilesets/agent_detect.py create mode 100644 tests/test_agent_detect.py diff --git a/mapbox_tilesets/agent_detect.py b/mapbox_tilesets/agent_detect.py new file mode 100644 index 0000000..1a8c432 --- /dev/null +++ b/mapbox_tilesets/agent_detect.py @@ -0,0 +1,68 @@ +"""Detect the AI coding agent (if any) driving this CLI invocation.""" + +import os + +# (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, value not inspected); 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 var in os.environ: + 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: + return value + + 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..a527b51 --- /dev/null +++ b/tests/test_agent_detect.py @@ -0,0 +1,81 @@ +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" 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) From 5af02d60874bf0f8708a61b87d9e3034ea90e7c2 Mon Sep 17 00:00:00 2001 From: ctufts Date: Fri, 17 Jul 2026 17:05:25 -0400 Subject: [PATCH 2/2] Sanitize agent env-var values before self-review findings ship - Reject fallback (AI_AGENT/AGENT) values outside a safe header charset so a stray newline/colon can no longer crash every CLI command with requests.exceptions.InvalidHeader. - Treat empty/whitespace-only harness env vars as unset, matching the fallback path's existing empty-value handling. --- mapbox_tilesets/agent_detect.py | 13 ++++++++++--- tests/test_agent_detect.py | 23 +++++++++++++++++++++++ 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/mapbox_tilesets/agent_detect.py b/mapbox_tilesets/agent_detect.py index 1a8c432..ed98155 100644 --- a/mapbox_tilesets/agent_detect.py +++ b/mapbox_tilesets/agent_detect.py @@ -1,10 +1,17 @@ """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, value not inspected); otherwise an exact-equality 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)]), @@ -55,14 +62,14 @@ def detect_agent(): for agent_id, conditions in _ALLOWLIST: for var, expected in conditions: if expected is None: - if var in os.environ: + 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: + if value and _SAFE_FALLBACK_ID.match(value): return value return None diff --git a/tests/test_agent_detect.py b/tests/test_agent_detect.py index a527b51..a37d268 100644 --- a/tests/test_agent_detect.py +++ b/tests/test_agent_detect.py @@ -79,3 +79,26 @@ 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