feat: detect all coding agents in user-agent, not just Claude Code

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, "").strip():
+                return agent_id
+
+    for env in STANDARD_ENV_VARS:
+        value = _sanitize(os.getenv(env, ""))
+        if value:
+            return value
+
+    return ""
+
+
+def suffix():
+    """Returns the " (via <id>)" 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 ""

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_suffix = agent.suffix()
+    if agent_suffix:
+        ua_components.append(agent_suffix.strip())
 
     user_agent = " ".join(ua_components)
     return user_agent

tests/test_agent.py

@@ -0,0 +1,124 @@
+"""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_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()
+        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()

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)