From 499c77c96edc15d1992d490f9b8042769711cb0d Mon Sep 17 00:00:00 2001 From: Luke Piette Date: Tue, 23 Jun 2026 11:15:47 -0400 Subject: [PATCH 1/2] feat: detect all coding agents in user-agent, not just Claude Code MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extends agent source tracking beyond CLAUDECODE to the full set of coding agent harnesses, mirroring runpodctl PR #280 and Hugging Face's public agent-harnesses registry so traffic is attributed under the same identifiers across surfaces. - Add runpod/agent.py: ordered harness registry (claude-code, cursor, cursor-cli, codex, gemini-cli, github-copilot, cline, replit, zed, etc.) plus a sanitized generic AI_AGENT fallback. - Wire agent.detect() into construct_user_agent(), replacing the single CLAUDECODE check. Emits "(via )" as before. - Tests for detection precedence, sanitization, and the user-agent string. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude --- runpod/agent.py | 99 +++++++++++++++++++++++++++++++++ runpod/user_agent.py | 6 +- tests/test_agent.py | 117 +++++++++++++++++++++++++++++++++++++++ tests/test_user_agent.py | 42 +++++++++++--- 4 files changed, 253 insertions(+), 11 deletions(-) create mode 100644 runpod/agent.py create mode 100644 tests/test_agent.py diff --git a/runpod/agent.py b/runpod/agent.py new file mode 100644 index 00000000..955031f2 --- /dev/null +++ b/runpod/agent.py @@ -0,0 +1,99 @@ +"""Detects which AI coding agent (if any) is driving the SDK. + +Detection is based on the environment variables that agent harnesses set in +the processes they spawn. The registry mirrors Hugging Face's public +agent-harnesses list so that traffic is attributed under the same agent +identifiers across tools: +https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/src/agent-harnesses.ts +""" + +import os + +# Each entry maps an agent identifier to the environment variables that +# identify it. Detection matches if ANY of the listed variables is set to a +# non-empty value. The list is checked in order and the first match wins; +# order matters so that more specific signals come before broader ones they +# can co-occur with (e.g. cowork before claude-code, cursor-cli before cursor). +HARNESSES = [ + ("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"]), + ("zed", ["ZED_TERM"]), + ("cursor-cli", ["CURSOR_AGENT"]), + ("cursor", ["CURSOR_TRACE_ID"]), +] + +# Generic variables any tool can set to identify itself. When set, the value is +# sanitized and used as the agent id. Only AI_AGENT is honored: a bare AGENT is +# too common in unrelated environments (CI runners, shell setups) and would +# attribute traffic to arbitrary values. +STANDARD_ENV_VARS = ["AI_AGENT"] + + +def known_env_vars(): + """Returns every environment variable the registry inspects, including the + standard AI_AGENT signal. Useful for tests that need to isolate detection + from the ambient environment. + """ + variables = [] + for _, env_vars in HARNESSES: + variables.extend(env_vars) + return variables + list(STANDARD_ENV_VARS) + + +def _sanitize(value): + """Keeps only User-Agent-safe characters ([A-Za-z0-9._-]), capped at 64 + characters, so an arbitrary env value cannot produce a malformed header. + """ + value = value.strip() + safe = [] + for char in value: + if len(safe) >= 64: + break + if char.isascii() and (char.isalnum() or char in "._-"): + safe.append(char) + return "".join(safe) + + +def detect(): + """Returns the identifier of the AI coding agent driving the SDK, or an + empty string if none is detected. Specific harness markers take priority + over the generic AI_AGENT signal. + """ + for agent_id, env_vars in HARNESSES: + for env in env_vars: + if os.getenv(env): + return agent_id + + for env in STANDARD_ENV_VARS: + value = _sanitize(os.getenv(env, "")) + if value: + return value + + return "" + + +def suffix(): + """Returns the " (via )" User-Agent fragment for the detected agent, or + an empty string when none is detected. Centralizing the fragment here keeps + the tag format identical across every client's User-Agent. + """ + agent_id = detect() + if agent_id: + return f" (via {agent_id})" + return "" diff --git a/runpod/user_agent.py b/runpod/user_agent.py index 88b013e4..6431de4e 100644 --- a/runpod/user_agent.py +++ b/runpod/user_agent.py @@ -3,6 +3,7 @@ import os import platform +from runpod import agent from runpod.version import __version__ as runpod_version @@ -25,8 +26,9 @@ def construct_user_agent(): if integration_method: ua_components.append(f"Integration/{integration_method}") - if os.getenv("CLAUDECODE") == "1": - ua_components.append("(via claude-code)") + agent_id = agent.detect() + if agent_id: + ua_components.append(f"(via {agent_id})") user_agent = " ".join(ua_components) return user_agent diff --git a/tests/test_agent.py b/tests/test_agent.py new file mode 100644 index 00000000..1202b426 --- /dev/null +++ b/tests/test_agent.py @@ -0,0 +1,117 @@ +"""Tests for the agent detection module.""" + +import os +import unittest +from unittest.mock import patch + +from runpod import agent + + +def _clean_env(): + """Returns an environment dict with every known agent signal removed.""" + env = dict(os.environ) + for key in agent.known_env_vars(): + env.pop(key, None) + return env + + +class TestDetect(unittest.TestCase): + """Test the agent.detect function.""" + + def test_no_agent(self): + """No agent env vars set means no detection.""" + with patch.dict(os.environ, _clean_env(), clear=True): + self.assertEqual(agent.detect(), "") + self.assertEqual(agent.suffix(), "") + + def test_claude_code(self): + """CLAUDECODE detects claude-code.""" + env = _clean_env() + env["CLAUDECODE"] = "1" + with patch.dict(os.environ, env, clear=True): + self.assertEqual(agent.detect(), "claude-code") + self.assertEqual(agent.suffix(), " (via claude-code)") + + def test_claude_code_alternate_var(self): + """CLAUDE_CODE also detects claude-code.""" + env = _clean_env() + env["CLAUDE_CODE"] = "1" + with patch.dict(os.environ, env, clear=True): + self.assertEqual(agent.detect(), "claude-code") + + def test_cursor(self): + """CURSOR_TRACE_ID detects cursor.""" + env = _clean_env() + env["CURSOR_TRACE_ID"] = "abc123" + with patch.dict(os.environ, env, clear=True): + self.assertEqual(agent.detect(), "cursor") + + def test_cursor_cli_before_cursor(self): + """cursor-cli is more specific and wins over cursor when both are set.""" + env = _clean_env() + env["CURSOR_AGENT"] = "1" + env["CURSOR_TRACE_ID"] = "abc123" + with patch.dict(os.environ, env, clear=True): + self.assertEqual(agent.detect(), "cursor-cli") + + def test_cowork_before_claude_code(self): + """cowork is more specific and wins over claude-code when both are set.""" + env = _clean_env() + env["CLAUDE_CODE_IS_COWORK"] = "1" + env["CLAUDECODE"] = "1" + with patch.dict(os.environ, env, clear=True): + self.assertEqual(agent.detect(), "cowork") + + def test_codex(self): + """A Codex marker detects codex.""" + env = _clean_env() + env["CODEX_THREAD_ID"] = "t-1" + with patch.dict(os.environ, env, clear=True): + self.assertEqual(agent.detect(), "codex") + + def test_gemini_cli(self): + """GEMINI_CLI detects gemini-cli.""" + env = _clean_env() + env["GEMINI_CLI"] = "1" + with patch.dict(os.environ, env, clear=True): + self.assertEqual(agent.detect(), "gemini-cli") + + def test_empty_value_not_detected(self): + """An env var set to empty string does not count as detection.""" + env = _clean_env() + env["CLAUDECODE"] = "" + with patch.dict(os.environ, env, clear=True): + self.assertEqual(agent.detect(), "") + + def test_ai_agent_generic_fallback(self): + """The generic AI_AGENT signal is used when no harness matches.""" + env = _clean_env() + env["AI_AGENT"] = "my-custom-agent" + with patch.dict(os.environ, env, clear=True): + self.assertEqual(agent.detect(), "my-custom-agent") + + def test_harness_wins_over_ai_agent(self): + """A specific harness marker takes priority over the generic AI_AGENT.""" + env = _clean_env() + env["AI_AGENT"] = "my-custom-agent" + env["CLAUDECODE"] = "1" + with patch.dict(os.environ, env, clear=True): + self.assertEqual(agent.detect(), "claude-code") + + def test_ai_agent_sanitized(self): + """AI_AGENT values are sanitized to User-Agent-safe characters.""" + env = _clean_env() + env["AI_AGENT"] = "bad value (with) chars!" + with patch.dict(os.environ, env, clear=True): + self.assertEqual(agent.detect(), "badvaluewithchars") + + def test_ai_agent_length_capped(self): + """AI_AGENT values are capped at 64 characters.""" + env = _clean_env() + env["AI_AGENT"] = "a" * 200 + with patch.dict(os.environ, env, clear=True): + self.assertEqual(agent.detect(), "a" * 64) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_user_agent.py b/tests/test_user_agent.py index d8678023..354d2a1c 100644 --- a/tests/test_user_agent.py +++ b/tests/test_user_agent.py @@ -5,9 +5,15 @@ from unittest.mock import patch from runpod import __version__ as runpod_version +from runpod import agent from runpod.user_agent import construct_user_agent +def _agent_env_keys(): + """Every env var that could inject an agent tag, plus the integration var.""" + return agent.known_env_vars() + ["RUNPOD_UA_INTEGRATION"] + + class TestConstructUserAgent(unittest.TestCase): """Test the construct_user_agent function.""" @@ -19,7 +25,7 @@ def test_user_agent_without_integration( self, mock_python_version, mock_machine, mock_release, mock_system ): """Test the User-Agent string without specifying an integration method.""" - saved = {k: os.environ.pop(k) for k in ("RUNPOD_UA_INTEGRATION", "CLAUDECODE") if k in os.environ} + saved = {k: os.environ.pop(k) for k in _agent_env_keys() if k in os.environ} expected_ua = f"RunPod-Python-SDK/{runpod_version} (Windows 10; AMD64) Language/Python 3.8.10" # pylint: disable=line-too-long self.assertEqual(construct_user_agent(), expected_ua) @@ -39,15 +45,14 @@ def test_user_agent_with_integration( self, mock_python_version, mock_machine, mock_release, mock_system ): """Test the User-Agent string with an integration method specified.""" - saved_claude = os.environ.pop("CLAUDECODE", None) + saved = {k: os.environ.pop(k) for k in _agent_env_keys() if k in os.environ} os.environ["RUNPOD_UA_INTEGRATION"] = "SkyPilot" expected_ua = f"RunPod-Python-SDK/{runpod_version} (Linux 5.4; x86_64) Language/Python 3.9.5 Integration/SkyPilot" # pylint: disable=line-too-long self.assertEqual(construct_user_agent(), expected_ua) os.environ.pop("RUNPOD_UA_INTEGRATION") - if saved_claude is not None: - os.environ["CLAUDECODE"] = saved_claude + os.environ.update(saved) assert mock_python_version.called assert mock_machine.called @@ -62,13 +67,15 @@ def test_user_agent_with_integration( def test_user_agent_with_claude_code( self, mock_python_version, mock_machine, mock_release, mock_system ): - """Test the User-Agent string includes claude-code agent tag.""" - os.environ.pop("RUNPOD_UA_INTEGRATION", None) + """Test the User-Agent string includes the claude-code agent tag.""" + saved = {k: os.environ.pop(k) for k in _agent_env_keys() if k in os.environ} os.environ["CLAUDECODE"] = "1" expected_ua = f"RunPod-Python-SDK/{runpod_version} (Linux 5.4; x86_64) Language/Python 3.9.5 (via claude-code)" self.assertEqual(construct_user_agent(), expected_ua) + os.environ.pop("CLAUDECODE", None) + os.environ.update(saved) @patch("runpod.user_agent.platform.system", return_value="Linux") @patch("runpod.user_agent.platform.release", return_value="5.4") @@ -77,12 +84,29 @@ def test_user_agent_with_claude_code( def test_user_agent_without_claude_code( self, mock_python_version, mock_machine, mock_release, mock_system ): - """Test the User-Agent string excludes agent tag when env var is not set.""" - saved = {k: os.environ.pop(k) for k in ("RUNPOD_UA_INTEGRATION", "CLAUDECODE") if k in os.environ} + """Test the User-Agent string excludes agent tag when no env var is set.""" + saved = {k: os.environ.pop(k) for k in _agent_env_keys() if k in os.environ} + + ua = construct_user_agent() + self.assertNotIn("(via ", ua) + + os.environ.update(saved) + + @patch("runpod.user_agent.platform.system", return_value="Linux") + @patch("runpod.user_agent.platform.release", return_value="5.4") + @patch("runpod.user_agent.platform.machine", return_value="x86_64") + @patch("runpod.user_agent.platform.python_version", return_value="3.9.5") + def test_user_agent_with_other_agent( + self, mock_python_version, mock_machine, mock_release, mock_system + ): + """Test the User-Agent string includes a non-Claude agent tag (e.g. cursor).""" + saved = {k: os.environ.pop(k) for k in _agent_env_keys() if k in os.environ} + os.environ["CURSOR_TRACE_ID"] = "abc123" ua = construct_user_agent() - self.assertNotIn("via claude-code", ua) + self.assertIn("(via cursor)", ua) + os.environ.pop("CURSOR_TRACE_ID", None) os.environ.update(saved) From 2e9dcd89a63f1283af7bcc382c6a2433fdf85806 Mon Sep 17 00:00:00 2001 From: Luke Piette Date: Tue, 23 Jun 2026 13:17:28 -0400 Subject: [PATCH 2/2] refactor: address Copilot review on agent tracking MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - user_agent.py: use agent.suffix() instead of re-implementing the "(via )" fragment, so the format lives in one place and cannot drift. - agent.detect(): strip env var values before matching, so a whitespace-only value does not count as detection (aligns harness matching with the docstring and the AI_AGENT path, which already strips). - Add a test for the whitespace-only case. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude --- runpod/agent.py | 2 +- runpod/user_agent.py | 6 +++--- tests/test_agent.py | 7 +++++++ 3 files changed, 11 insertions(+), 4 deletions(-) diff --git a/runpod/agent.py b/runpod/agent.py index 955031f2..619eb469 100644 --- a/runpod/agent.py +++ b/runpod/agent.py @@ -77,7 +77,7 @@ def detect(): """ for agent_id, env_vars in HARNESSES: for env in env_vars: - if os.getenv(env): + if os.getenv(env, "").strip(): return agent_id for env in STANDARD_ENV_VARS: diff --git a/runpod/user_agent.py b/runpod/user_agent.py index 6431de4e..021745de 100644 --- a/runpod/user_agent.py +++ b/runpod/user_agent.py @@ -26,9 +26,9 @@ def construct_user_agent(): if integration_method: ua_components.append(f"Integration/{integration_method}") - agent_id = agent.detect() - if agent_id: - ua_components.append(f"(via {agent_id})") + agent_suffix = agent.suffix() + if agent_suffix: + ua_components.append(agent_suffix.strip()) user_agent = " ".join(ua_components) return user_agent diff --git a/tests/test_agent.py b/tests/test_agent.py index 1202b426..ea9dadf2 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -83,6 +83,13 @@ def test_empty_value_not_detected(self): with patch.dict(os.environ, env, clear=True): self.assertEqual(agent.detect(), "") + def test_whitespace_value_not_detected(self): + """An env var set to whitespace only does not count as detection.""" + env = _clean_env() + env["CLAUDECODE"] = " " + with patch.dict(os.environ, env, clear=True): + self.assertEqual(agent.detect(), "") + def test_ai_agent_generic_fallback(self): """The generic AI_AGENT signal is used when no harness matches.""" env = _clean_env()