PYTHON-5929 Add coding agent env var to handshake metadata

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
 .........

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)

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})