diff --git a/doc/changelog.rst b/doc/changelog.rst index fd8ef94334..94d387345c 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -43,6 +43,13 @@ PyMongo 4.19 brings a number of changes including: security considerations. - ``srvAllowedHostsSuffix`` may now be set to a single label reserved for private or special use, such as ``localhost``, ``test``, or ``internal``. +- The client handshake metadata now reports a detected coding agent in + ``client.env.agent``. The driver checks a fixed list of agent environment + variables (for example ``CLAUDECODE``, ``CURSOR_AGENT`` and + ``CODEX_SANDBOX``) and reports a fixed name for the first one that is set. + If none is set, the driver uses the generic ``AI_AGENT`` variable: ``1`` or + ``true`` is reported as ``ai_agent``, and any other value is trimmed, + lowercased and truncated to 64 bytes. Bug fixes ......... diff --git a/pymongo/pool_options.py b/pymongo/pool_options.py index f20033ad0e..a80d17a9ef 100644 --- a/pymongo/pool_options.py +++ b/pymongo/pool_options.py @@ -151,6 +151,46 @@ def _is_faas() -> bool: return _is_lambda() or _is_azure_func() or _is_gcp_func() or _is_vercel() +# Known coding agents, checked in order. The first populated variable gives +# client.env.agent, whatever its value. Matches mongosh detection. +# See DRIVERS-3529 and PYTHON-5929. +_AGENT_ENV_VARS = [ + ("CLAUDECODE", "claude_code"), + ("CLAUDE_CODE_ENTRYPOINT", "claude_code"), + ("CURSOR_AGENT", "cursor"), + ("CODEX_SANDBOX", "codex_cli"), + ("CLINE_ACTIVE", "cline"), + ("GEMINI_CLI", "gemini_cli"), + ("AUGMENT_AGENT", "auggie_cli"), + ("OPENCODE_CLIENT", "opencode_client"), + ("TRAE_AI_SHELL_ID", "trae_ai"), + ("GOOSE_TERMINAL", "goose"), + ("GOOSE_AGENT", "goose"), +] + +# Generic agent variable, checked last so a known agent keeps its fixed name. +_GENERIC_AGENT_ENV_VAR = "AI_AGENT" + +# Maximum size in bytes of a normalized AI_AGENT value. +_MAX_AGENT_SIZE = 64 + + +def _metadata_agent() -> Optional[str]: + """Detect a coding agent from the environment for client.env.agent.""" + for var, name in _AGENT_ENV_VARS: + # Unset, empty or whitespace-only is not populated. + if (os.getenv(var) or "").strip(): + return name + agent = (os.getenv(_GENERIC_AGENT_ENV_VAR) or "").strip().lower() + if not agent: + return None + if agent in ("1", "true"): + return "ai_agent" + # Largest valid UTF-8 prefix of _MAX_AGENT_SIZE bytes. "ignore" drops a + # split character instead of replacing it with U+FFFD. + return agent.encode()[:_MAX_AGENT_SIZE].decode(errors="ignore") + + def _getenv_int(key: str) -> Optional[int]: """Like os.getenv but returns an int, or None if the value is missing/malformed.""" val = os.getenv(key) @@ -167,6 +207,9 @@ def _metadata_env() -> dict[str, Any]: container = get_container_env_info() if container: env["container"] = container + agent = _metadata_agent() + if agent: + env["agent"] = agent # Skip if multiple (or no) envs are matched. if (_is_lambda(), _is_azure_func(), _is_gcp_func(), _is_vercel()).count(True) != 1: return env @@ -236,15 +279,27 @@ def _truncate_metadata(metadata: MutableMapping[str, Any]) -> None: size = len(bson.encode(metadata)) if size <= _MAX_METADATA_SIZE: return - # 1. Omit fields from env except env.name. - env_name = metadata.get("env", {}).get("name") - if env_name: - env = {"name": env_name} - size += _element_size("env", env) - _element_size("env", metadata["env"]) - metadata["env"] = env + # 1. Omit fields from env except env.name and env.agent. + env = metadata.get("env", {}) + trimmed_env = {k: env[k] for k in ("name", "agent") if k in env} + if trimmed_env: + size += _element_size("env", trimmed_env) - _element_size("env", env) + metadata["env"] = trimmed_env + if size <= _MAX_METADATA_SIZE: + return + # 2. Omit env.agent. It goes before env.name, which drivers have reported + # for longer. If env has no remaining fields, omit env entirely. + env = metadata.get("env") + if env is not None and "agent" in env: + new_env = {k: v for k, v in env.items() if k != "agent"} + size += _element_size("env", new_env) - _element_size("env", env) + if new_env: + metadata["env"] = new_env + else: + del metadata["env"] if size <= _MAX_METADATA_SIZE: return - # 2. Omit fields from os except os.type. + # 3. Omit fields from os except os.type. os_type = metadata.get("os", {}).get("type") if os_type: old_os = metadata["os"] @@ -253,13 +308,13 @@ def _truncate_metadata(metadata: MutableMapping[str, Any]) -> None: metadata["os"] = new_os if size <= _MAX_METADATA_SIZE: return - # 3. Omit the env document entirely. + # 4. Omit the env document entirely. env = metadata.pop("env", None) if env is not None: size -= _element_size("env", env) if size <= _MAX_METADATA_SIZE: return - # 4. Truncate platform. + # 5. Truncate platform. overflow = size - _MAX_METADATA_SIZE plat = metadata.get("platform", "") if plat: @@ -277,7 +332,7 @@ def _truncate_metadata(metadata: MutableMapping[str, Any]) -> None: plat = metadata.pop("platform", None) if plat is not None: size -= _element_size("platform", plat) - # 5. Truncate driver info, keeping name and version 1:1 index-aligned. + # 6. Truncate driver info, keeping name and version 1:1 index-aligned. driver = metadata.get("driver", {}) if driver: # Keep the name and version segments paired so they stay 1:1 aligned, @@ -298,7 +353,7 @@ def _truncate_metadata(metadata: MutableMapping[str, Any]) -> None: overflow = size - _MAX_METADATA_SIZE while overflow > 0 and len(pairs) > 1: # A single remaining pair never exceeds the limit: it is the base - # driver pair, which steps 1-4 left small enough to fit. + # driver pair, which steps 1-5 left small enough to fit. last_name, last_version = pairs[-1] if last_version: new_version = _truncate_utf8(last_version, overflow) diff --git a/test/asynchronous/test_client.py b/test/asynchronous/test_client.py index 46249f762f..bb7d893733 100644 --- a/test/asynchronous/test_client.py +++ b/test/asynchronous/test_client.py @@ -91,6 +91,7 @@ ) from pymongo.monitoring import ServerHeartbeatListener, ServerHeartbeatStartedEvent from pymongo.pool_options import ( + _MAX_AGENT_SIZE, _MAX_METADATA_SIZE, _METADATA, ENV_VAR_K8S, @@ -136,6 +137,7 @@ gevent_monkey_patched, is_greenthread_patched, lazy_client_trial, + no_ambient_agent, one, suppress_fork_deprecation, ) @@ -389,6 +391,7 @@ async def test_read_preference(self): ) self.assertEqual(c.read_preference, ReadPreference.NEAREST) + @no_ambient_agent() async def test_metadata(self): metadata = copy.deepcopy(_METADATA) if has_c(): @@ -513,6 +516,42 @@ def metadata_for(name_len: int) -> dict[str, Any]: {"name": "PyMongo|" + "W" * name_len, "version": "1.0|1.0"}, ) + async def test_metadata_truncation_omits_agent_before_name(self): + # Truncation keeps env.name and env.agent first, then omits env.agent + # alone. An env holding only agent has no remaining field, so the env + # document is omitted entirely. + + def metadata_for(plat_len: int, with_name: bool) -> dict[str, Any]: + env: dict[str, Any] = {"agent": "claude_code"} + if with_name: + env["name"] = "azure.func" + return {"env": env, "platform": "p" * plat_len} + + # Size platform so the document is one byte over the limit, less than + # the bytes the agent element costs: dropping agent alone must make it + # fit without touching platform. + plat_len = next( + n + for n in range(1, 2 * _MAX_METADATA_SIZE) + if len(bson.encode(metadata_for(n, True))) == _MAX_METADATA_SIZE + 1 + ) + metadata = metadata_for(plat_len, True) + _truncate_metadata(metadata) + self.assertLessEqual(len(bson.encode(metadata)), _MAX_METADATA_SIZE) + self.assertEqual(metadata["env"], {"name": "azure.func"}) + self.assertEqual(metadata["platform"], "p" * plat_len) + + plat_len = next( + n + for n in range(1, 2 * _MAX_METADATA_SIZE) + if len(bson.encode(metadata_for(n, False))) == _MAX_METADATA_SIZE + 1 + ) + metadata = metadata_for(plat_len, False) + _truncate_metadata(metadata) + self.assertLessEqual(len(bson.encode(metadata)), _MAX_METADATA_SIZE) + self.assertNotIn("env", metadata) + self.assertEqual(metadata["platform"], "p" * plat_len) + async def test_metadata_append_is_bounded(self): # Successive appends must stay within the limit and keep name and # version index-aligned after truncation. Once the metadata saturates, @@ -567,6 +606,7 @@ async def test_metadata_deduplicates_none_and_empty(self): self.assertEqual(metadata["driver"]["name"], names) self.assertEqual(metadata["driver"]["version"], vers) + @no_ambient_agent() @mock.patch.dict("os.environ", {ENV_VAR_K8S: "1"}) def test_container_metadata(self): metadata = copy.deepcopy(_METADATA) @@ -2332,7 +2372,9 @@ def test_sigstop_sigcont(self): self.assertNotIn("ServerHeartbeatFailedEvent", log_output) async def _test_handshake(self, env_vars, expected_env): - with patch.dict("os.environ", env_vars): + # Clear ambient agent vars (e.g. AI_AGENT set by the CI runner) so + # detection only reflects env_vars. + with no_ambient_agent(keep=env_vars), patch.dict("os.environ", env_vars): metadata = copy.deepcopy(_METADATA) if has_c(): metadata["driver"]["name"] = "PyMongo|c|async" @@ -2440,6 +2482,59 @@ async def test_handshake_09_container_with_provider(self): }, ) + async def test_handshake_10_agent_known(self): + # A known agent env var maps to its fixed name, regardless of value. + await self._test_handshake({"CLAUDECODE": "1"}, {"agent": "claude_code"}) + await self._test_handshake({"CURSOR_AGENT": "some-value-42"}, {"agent": "cursor"}) + await self._test_handshake({"OPENCODE_CLIENT": "1"}, {"agent": "opencode_client"}) + + async def test_handshake_10b_agent_known_precedence(self): + # The first var in _AGENT_ENV_VARS order wins, not the first in the + # environment dict. + await self._test_handshake({"GEMINI_CLI": "1", "CURSOR_AGENT": "1"}, {"agent": "cursor"}) + + async def test_handshake_11_agent_known_beats_generic(self): + # A known agent wins over AI_AGENT, so a versioned AI_AGENT value + # cannot mask it. + await self._test_handshake( + {"AI_AGENT": "custom-agent", "CLAUDECODE": "1"}, {"agent": "claude_code"} + ) + + async def test_handshake_12_agent_generic(self): + # A descriptive value is used as-is. "1" and "true" map to "ai_agent". + await self._test_handshake({"AI_AGENT": "custom-agent"}, {"agent": "custom-agent"}) + await self._test_handshake({"AI_AGENT": "1"}, {"agent": "ai_agent"}) + await self._test_handshake({"AI_AGENT": "true"}, {"agent": "ai_agent"}) + + async def test_handshake_13_agent_generic_normalized(self): + # AI_AGENT is trimmed and lowercased. + await self._test_handshake( + {"AI_AGENT": " Claude-Code_2-1-238_Agent "}, {"agent": "claude-code_2-1-238_agent"} + ) + + async def test_handshake_14_agent_generic_truncated(self): + # A long value is truncated to _MAX_AGENT_SIZE bytes. + await self._test_handshake({"AI_AGENT": "a" * 100}, {"agent": "a" * _MAX_AGENT_SIZE}) + + async def test_handshake_14b_agent_generic_truncated_on_boundary(self): + # The byte limit falls inside the two-byte "é", so the character is + # dropped. No part of it, and no U+FFFD, may appear. + value = "a" * (_MAX_AGENT_SIZE - 1) + "é" + await self._test_handshake({"AI_AGENT": value}, {"agent": "a" * (_MAX_AGENT_SIZE - 1)}) + + async def test_handshake_15_agent_unset(self): + # An empty or whitespace-only value counts as unset. + await self._test_handshake({"AI_AGENT": ""}, None) + await self._test_handshake({"AI_AGENT": " "}, None) + await self._test_handshake({"CLAUDECODE": " "}, None) + + async def test_handshake_16_agent_with_provider(self): + # agent is reported alongside a FaaS provider. + await self._test_handshake( + {"FUNCTIONS_WORKER_RUNTIME": "python", "CLAUDECODE": "1"}, + {"name": "azure.func", "agent": "claude_code"}, + ) + def test_dict_hints(self): self.db.coll.find(hint={"x": 1}) diff --git a/test/test_client.py b/test/test_client.py index a2640def28..6491b2c97a 100644 --- a/test/test_client.py +++ b/test/test_client.py @@ -82,6 +82,7 @@ ) from pymongo.monitoring import ServerHeartbeatListener, ServerHeartbeatStartedEvent from pymongo.pool_options import ( + _MAX_AGENT_SIZE, _MAX_METADATA_SIZE, _METADATA, ENV_VAR_K8S, @@ -135,6 +136,7 @@ gevent_monkey_patched, is_greenthread_patched, lazy_client_trial, + no_ambient_agent, one, suppress_fork_deprecation, ) @@ -382,6 +384,7 @@ def test_read_preference(self): ) self.assertEqual(c.read_preference, ReadPreference.NEAREST) + @no_ambient_agent() def test_metadata(self): metadata = copy.deepcopy(_METADATA) if has_c(): @@ -506,6 +509,42 @@ def metadata_for(name_len: int) -> dict[str, Any]: {"name": "PyMongo|" + "W" * name_len, "version": "1.0|1.0"}, ) + def test_metadata_truncation_omits_agent_before_name(self): + # Truncation keeps env.name and env.agent first, then omits env.agent + # alone. An env holding only agent has no remaining field, so the env + # document is omitted entirely. + + def metadata_for(plat_len: int, with_name: bool) -> dict[str, Any]: + env: dict[str, Any] = {"agent": "claude_code"} + if with_name: + env["name"] = "azure.func" + return {"env": env, "platform": "p" * plat_len} + + # Size platform so the document is one byte over the limit, less than + # the bytes the agent element costs: dropping agent alone must make it + # fit without touching platform. + plat_len = next( + n + for n in range(1, 2 * _MAX_METADATA_SIZE) + if len(bson.encode(metadata_for(n, True))) == _MAX_METADATA_SIZE + 1 + ) + metadata = metadata_for(plat_len, True) + _truncate_metadata(metadata) + self.assertLessEqual(len(bson.encode(metadata)), _MAX_METADATA_SIZE) + self.assertEqual(metadata["env"], {"name": "azure.func"}) + self.assertEqual(metadata["platform"], "p" * plat_len) + + plat_len = next( + n + for n in range(1, 2 * _MAX_METADATA_SIZE) + if len(bson.encode(metadata_for(n, False))) == _MAX_METADATA_SIZE + 1 + ) + metadata = metadata_for(plat_len, False) + _truncate_metadata(metadata) + self.assertLessEqual(len(bson.encode(metadata)), _MAX_METADATA_SIZE) + self.assertNotIn("env", metadata) + self.assertEqual(metadata["platform"], "p" * plat_len) + def test_metadata_append_is_bounded(self): # Successive appends must stay within the limit and keep name and # version index-aligned after truncation. Once the metadata saturates, @@ -560,6 +599,7 @@ def test_metadata_deduplicates_none_and_empty(self): self.assertEqual(metadata["driver"]["name"], names) self.assertEqual(metadata["driver"]["version"], vers) + @no_ambient_agent() @mock.patch.dict("os.environ", {ENV_VAR_K8S: "1"}) def test_container_metadata(self): metadata = copy.deepcopy(_METADATA) @@ -2285,7 +2325,9 @@ def test_sigstop_sigcont(self): self.assertNotIn("ServerHeartbeatFailedEvent", log_output) def _test_handshake(self, env_vars, expected_env): - with patch.dict("os.environ", env_vars): + # Clear ambient agent vars (e.g. AI_AGENT set by the CI runner) so + # detection only reflects env_vars. + with no_ambient_agent(keep=env_vars), patch.dict("os.environ", env_vars): metadata = copy.deepcopy(_METADATA) if has_c(): metadata["driver"]["name"] = "PyMongo|c" @@ -2393,6 +2435,59 @@ def test_handshake_09_container_with_provider(self): }, ) + def test_handshake_10_agent_known(self): + # A known agent env var maps to its fixed name, regardless of value. + self._test_handshake({"CLAUDECODE": "1"}, {"agent": "claude_code"}) + self._test_handshake({"CURSOR_AGENT": "some-value-42"}, {"agent": "cursor"}) + self._test_handshake({"OPENCODE_CLIENT": "1"}, {"agent": "opencode_client"}) + + def test_handshake_10b_agent_known_precedence(self): + # The first var in _AGENT_ENV_VARS order wins, not the first in the + # environment dict. + self._test_handshake({"GEMINI_CLI": "1", "CURSOR_AGENT": "1"}, {"agent": "cursor"}) + + def test_handshake_11_agent_known_beats_generic(self): + # A known agent wins over AI_AGENT, so a versioned AI_AGENT value + # cannot mask it. + self._test_handshake( + {"AI_AGENT": "custom-agent", "CLAUDECODE": "1"}, {"agent": "claude_code"} + ) + + def test_handshake_12_agent_generic(self): + # A descriptive value is used as-is. "1" and "true" map to "ai_agent". + self._test_handshake({"AI_AGENT": "custom-agent"}, {"agent": "custom-agent"}) + self._test_handshake({"AI_AGENT": "1"}, {"agent": "ai_agent"}) + self._test_handshake({"AI_AGENT": "true"}, {"agent": "ai_agent"}) + + def test_handshake_13_agent_generic_normalized(self): + # AI_AGENT is trimmed and lowercased. + self._test_handshake( + {"AI_AGENT": " Claude-Code_2-1-238_Agent "}, {"agent": "claude-code_2-1-238_agent"} + ) + + def test_handshake_14_agent_generic_truncated(self): + # A long value is truncated to _MAX_AGENT_SIZE bytes. + self._test_handshake({"AI_AGENT": "a" * 100}, {"agent": "a" * _MAX_AGENT_SIZE}) + + def test_handshake_14b_agent_generic_truncated_on_boundary(self): + # The byte limit falls inside the two-byte "é", so the character is + # dropped. No part of it, and no U+FFFD, may appear. + value = "a" * (_MAX_AGENT_SIZE - 1) + "é" + self._test_handshake({"AI_AGENT": value}, {"agent": "a" * (_MAX_AGENT_SIZE - 1)}) + + def test_handshake_15_agent_unset(self): + # An empty or whitespace-only value counts as unset. + self._test_handshake({"AI_AGENT": ""}, None) + self._test_handshake({"AI_AGENT": " "}, None) + self._test_handshake({"CLAUDECODE": " "}, None) + + def test_handshake_16_agent_with_provider(self): + # agent is reported alongside a FaaS provider. + self._test_handshake( + {"FUNCTIONS_WORKER_RUNTIME": "python", "CLAUDECODE": "1"}, + {"name": "azure.func", "agent": "claude_code"}, + ) + def test_dict_hints(self): self.db.coll.find(hint={"x": 1}) diff --git a/test/utils_shared.py b/test/utils_shared.py index fb209d69a0..2b4284f9c5 100644 --- a/test/utils_shared.py +++ b/test/utils_shared.py @@ -32,6 +32,7 @@ from functools import partial from inspect import iscoroutinefunction from typing import Any +from unittest.mock import patch from bson.objectid import ObjectId from pymongo import monitoring, operations, read_preferences @@ -53,7 +54,7 @@ PoolCreatedEvent, PoolReadyEvent, ) -from pymongo.pool_options import _METADATA +from pymongo.pool_options import _AGENT_ENV_VARS, _GENERIC_AGENT_ENV_VAR, _METADATA from pymongo.pool_shared import _CancellationContext, _PoolGeneration from pymongo.read_concern import ReadConcern from pymongo.server_type import SERVER_TYPE @@ -616,6 +617,20 @@ def suppress_fork_deprecation(): yield +def no_ambient_agent(keep=()): + """Clear the coding agent env vars the handshake reads. + + The test process itself often runs under a coding agent, which would add + client.env.agent to the handshake metadata. Vars named in `keep` are left + alone so a test can set its own (PYTHON-5929). + + :param keep: env var names to leave unchanged. + :return: a patch.dict context manager / decorator. + """ + agent_vars = [_GENERIC_AGENT_ENV_VAR, *(var for var, _ in _AGENT_ENV_VARS)] + return patch.dict("os.environ", {var: "" for var in agent_vars if var not in keep}) + + def parse_read_preference(pref): # Make first letter lowercase to match read_pref's modes. mode_string = pref.get("mode", "primary")