From c05c43717ac9d85e88b132c0c7bfb63915b54e85 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Tue, 26 May 2026 11:10:03 -0700 Subject: [PATCH 1/2] feat: add agent source tracking to user-agent header --- runpod/user_agent.py | 3 +++ tests/test_user_agent.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/runpod/user_agent.py b/runpod/user_agent.py index 714d43b6..88b013e4 100644 --- a/runpod/user_agent.py +++ b/runpod/user_agent.py @@ -25,6 +25,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)") + user_agent = " ".join(ua_components) return user_agent diff --git a/tests/test_user_agent.py b/tests/test_user_agent.py index 10dbf9c3..3d4fa03d 100644 --- a/tests/test_user_agent.py +++ b/tests/test_user_agent.py @@ -51,5 +51,36 @@ def test_user_agent_with_integration( assert mock_system.called + @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") + @patch.dict(os.environ, {"CLAUDECODE": "1"}, clear=False) + 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.""" + if "RUNPOD_UA_INTEGRATION" in os.environ: + del os.environ["RUNPOD_UA_INTEGRATION"] + + 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) + + @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_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.""" + for key in ("RUNPOD_UA_INTEGRATION", "CLAUDECODE"): + if key in os.environ: + del os.environ[key] + + ua = construct_user_agent() + self.assertNotIn("via claude-code", ua) + + if __name__ == "__main__": unittest.main() From 09882b1f5c9b6ad99285ffedf069f83a096a84d1 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Tue, 26 May 2026 11:46:00 -0700 Subject: [PATCH 2/2] fix: use explicit env var set/restore in user agent tests --- tests/test_user_agent.py | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/tests/test_user_agent.py b/tests/test_user_agent.py index 3d4fa03d..d8678023 100644 --- a/tests/test_user_agent.py +++ b/tests/test_user_agent.py @@ -19,12 +19,13 @@ 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.""" - if "RUNPOD_UA_INTEGRATION" in os.environ: - del os.environ["RUNPOD_UA_INTEGRATION"] + saved = {k: os.environ.pop(k) for k in ("RUNPOD_UA_INTEGRATION", "CLAUDECODE") 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) + os.environ.update(saved) + assert mock_python_version.called assert mock_machine.called assert mock_release.called @@ -34,16 +35,19 @@ def test_user_agent_without_integration( @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") - @patch.dict(os.environ, {"RUNPOD_UA_INTEGRATION": "SkyPilot"}) 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.""" - 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 - + saved_claude = os.environ.pop("CLAUDECODE", None) 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 assert mock_python_version.called assert mock_machine.called @@ -55,16 +59,16 @@ def test_user_agent_with_integration( @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") - @patch.dict(os.environ, {"CLAUDECODE": "1"}, clear=False) 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.""" - if "RUNPOD_UA_INTEGRATION" in os.environ: - del os.environ["RUNPOD_UA_INTEGRATION"] + os.environ.pop("RUNPOD_UA_INTEGRATION", None) + 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) @patch("runpod.user_agent.platform.system", return_value="Linux") @patch("runpod.user_agent.platform.release", return_value="5.4") @@ -74,13 +78,13 @@ 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.""" - for key in ("RUNPOD_UA_INTEGRATION", "CLAUDECODE"): - if key in os.environ: - del os.environ[key] + saved = {k: os.environ.pop(k) for k in ("RUNPOD_UA_INTEGRATION", "CLAUDECODE") if k in os.environ} ua = construct_user_agent() self.assertNotIn("via claude-code", ua) + os.environ.update(saved) + if __name__ == "__main__": unittest.main()