PYTHON-5929 Add coding agent env var to handshake metadata

doc/changelog.rst

@@ -4,6 +4,9 @@ Changelog
 Changes in Version 4.20.0 (2026/XX/XX)
 --------------------------------------
 
+- The client handshake metadata now reports a detected coding agent in
+  ``client.env.agent`` (`PYTHON-5929`_).
+
 Bug fixes
 .........
 
@@ -23,6 +26,7 @@ Bug fixes
   :meth:`~bson.objectid.ObjectId.is_valid` returns ``False`` for such a string
   (`PYTHON-5873`_).
 
+.. _PYTHON-5929: https://jira.mongodb.org/browse/PYTHON-5929
 .. _PYTHON-6172: https://jira.mongodb.org/browse/PYTHON-6172
 .. _PYTHON-6149: https://jira.mongodb.org/browse/PYTHON-6149
 .. _PYTHON-5873: https://jira.mongodb.org/browse/PYTHON-5873

pymongo/pool_options.py

@@ -142,6 +142,45 @@ def _is_faas() -> bool:
     return _is_lambda() or _is_azure_func() or _is_gcp_func() or _is_vercel()
 
 
+# Ordered per the handshake spec: the first populated variable
+# gives client.env.agent, whatever its value.
+_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", "opencode_client"),
+    ("OPENCODE_CLIENT", "opencode_client"),
+    ("TRAE_AI_SHELL_ID", "trae_ai"),
+    ("GOOSE_TERMINAL", "goose"),
+    ("GOOSE_AGENT", "goose"),
+]
+
+# Evaluated last so a known agent keeps its fixed name.
+_GENERIC_AGENT_ENV_VAR = "AI_AGENT"
+
+_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:
+        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
+    # character split by the byte limit and surrogate-escaped bytes from a
+    # malformed value so detection never raises.
+    return agent.encode(errors="ignore")[:_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)
@@ -158,6 +197,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
@@ -227,15 +269,29 @@ 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. env.agent is omitted before env.name because drivers
+    # have reported env.name since before env.agent existed.
+    env = metadata.get("env")
+    if env is not None and "agent" in env:
+        old_size = _element_size("env", env)
+        new_env = {k: v for k, v in env.items() if k != "agent"}
+        if new_env:
+            size += _element_size("env", new_env) - old_size
+            metadata["env"] = new_env
+        else:
+            size -= old_size
+            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"]
@@ -244,13 +300,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:
@@ -268,7 +324,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,
@@ -289,7 +345,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,10 +91,12 @@
 )
 from pymongo.monitoring import ServerHeartbeatListener, ServerHeartbeatStartedEvent
 from pymongo.pool_options import (
+    _MAX_AGENT_SIZE,
     _MAX_METADATA_SIZE,
     _METADATA,
     ENV_VAR_K8S,
     PoolOptions,
+    _element_size,
     _truncate_metadata,
 )
 from pymongo.read_preferences import ReadPreference
@@ -138,6 +140,7 @@
     lazy_client_trial,
     one,
     suppress_fork_deprecation,
+    unset_agent_env,
 )
 
 _IS_SYNC = False
@@ -389,6 +392,7 @@ async def test_read_preference(self):
         )
         self.assertEqual(c.read_preference, ReadPreference.NEAREST)
 
+    @unset_agent_env()
     async def test_metadata(self):
         metadata = copy.deepcopy(_METADATA)
         if has_c():
@@ -422,6 +426,7 @@ async def test_metadata_bad_driver_options(self):
         with self.assertRaises(TypeError):
             self.simple_client(driver=("Foo", "1", "a"))
 
+    @unset_agent_env()
     async def test_metadata_appends_driver_info(self):
         metadata = _metadata_with_appended_driver(_IS_SYNC, "FooDriver", "1.2.3")
         client = self.simple_client(
@@ -513,6 +518,72 @@ def metadata_for(name_len: int) -> dict[str, Any]:
             {"name": "PyMongo|" + "W" * name_len, "version": "1.0|1.0"},
         )
 
+    def _metadata_with_agent_env(self, with_name: bool) -> dict[str, Any]:
+        """Create metadata just over the size limit enforced by omitting `env.agent`."""
+        env: dict[str, Any] = {"agent": "claude_code"}
+        if with_name:
+            env["name"] = "azure.func"
+        os_meta = {"type": "Darwin", "name": "Darwin", "architecture": "arm64", "version": "26.7.1"}
+
+        def metadata_for(plat_len: int) -> dict[str, Any]:
+            return {"env": env, "os": os_meta, "platform": "p" * plat_len}
+
+        target = (
+            _MAX_METADATA_SIZE + 1
+            if with_name
+            else _MAX_METADATA_SIZE - 1 + _element_size("env", env)
+        )
+        # Each platform character adds exactly one byte.
+        plat_len = target - len(bson.encode(metadata_for(0)))
+        return metadata_for(plat_len)
+
+    async def test_metadata_truncation_keeps_name_and_agent(self):
+        # Dropping the other env fields alone must bring the document back
+        # under the limit.
+        env: dict[str, Any] = {
+            "name": "azure.func",
+            "agent": "claude_code",
+            "region": "us-east-2",
+            "memory_mb": 1024,
+        }
+        os_meta = {"type": "Darwin", "name": "Darwin", "architecture": "arm64", "version": "26.7.1"}
+
+        def metadata_for(plat_len: int) -> dict[str, Any]:
+            return {"env": env, "os": os_meta, "platform": "p" * plat_len}
+
+        # Each platform character adds exactly one byte.
+        plat_len = _MAX_METADATA_SIZE + 1 - len(bson.encode(metadata_for(0)))
+        metadata = metadata_for(plat_len)
+        plat = metadata["platform"]
+        _truncate_metadata(metadata)
+        self.assertLessEqual(len(bson.encode(metadata)), _MAX_METADATA_SIZE)
+        self.assertEqual(metadata["env"], {"name": "azure.func", "agent": "claude_code"})
+        self.assertEqual(metadata["os"], os_meta)
+        self.assertEqual(metadata["platform"], plat)
+
+    async def test_metadata_truncation_omits_agent_before_name(self):
+        metadata = self._metadata_with_agent_env(with_name=True)
+        os_meta = metadata["os"]
+        plat = metadata["platform"]
+        _truncate_metadata(metadata)
+        self.assertLessEqual(len(bson.encode(metadata)), _MAX_METADATA_SIZE)
+        self.assertEqual(metadata["env"], {"name": "azure.func"})
+        self.assertEqual(metadata["os"], os_meta)
+        self.assertEqual(metadata["platform"], plat)
+
+    async def test_metadata_truncation_omits_agent_only_env(self):
+        # env holds only agent so the whole env document is omitted. The size
+        # accounting must record the full env element so no later field is
+        # touched.
+        metadata = self._metadata_with_agent_env(with_name=False)
+        os_meta = metadata["os"]
+        plat = metadata["platform"]
+        _truncate_metadata(metadata)
+        self.assertLessEqual(len(bson.encode(metadata)), _MAX_METADATA_SIZE)
+        self.assertNotIn("env", metadata)
+        self.assertEqual(metadata["os"], os_meta)
+        self.assertEqual(metadata["platform"], plat)
+
     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 +638,7 @@ async def test_metadata_deduplicates_none_and_empty(self):
         self.assertEqual(metadata["driver"]["name"], names)
         self.assertEqual(metadata["driver"]["version"], vers)
 
+    @unset_agent_env()
     @mock.patch.dict("os.environ", {ENV_VAR_K8S: "1"})
     def test_container_metadata(self):
         metadata = copy.deepcopy(_METADATA)
@@ -2327,7 +2399,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):
+        # Unset agent env vars already set in the test process (e.g. AI_AGENT
+        # set by the CI runner) so detection only reflects env_vars.
+        with unset_agent_env(), patch.dict("os.environ", env_vars):
             metadata = copy.deepcopy(_METADATA)
             if has_c():
                 metadata["driver"]["name"] = "PyMongo|c|async"
@@ -2435,6 +2509,72 @@ async def test_handshake_09_container_with_provider(self):
             },
         )
 
+    async def test_handshake_10_agent_known(self):
+        await self._test_handshake({"CLAUDECODE": "1"}, {"agent": "claude_code"})
+
+    async def test_handshake_11_agent_known_fixed_name(self):
+        await self._test_handshake({"CURSOR_AGENT": "some-value-42"}, {"agent": "cursor"})
+
+    async def test_handshake_12_agent_known_opencode(self):
+        await self._test_handshake({"OPENCODE": "1"}, {"agent": "opencode_client"})
+
+    async def test_handshake_13_agent_known_opencode_client(self):
+        await self._test_handshake({"OPENCODE_CLIENT": "1"}, {"agent": "opencode_client"})
+
+    async def test_handshake_14_agent_known_precedence(self):
+        await self._test_handshake({"GEMINI_CLI": "1", "CURSOR_AGENT": "1"}, {"agent": "cursor"})
+
+    async def test_handshake_15_agent_known_beats_generic(self):
+        await self._test_handshake(
+            {"AI_AGENT": "custom-agent", "CLAUDECODE": "1"}, {"agent": "claude_code"}
+        )
+
+    async def test_handshake_16_agent_generic_descriptive(self):
+        await self._test_handshake({"AI_AGENT": "custom-agent"}, {"agent": "custom-agent"})
+
+    async def test_handshake_17_agent_generic_boolean(self):
+        await self._test_handshake({"AI_AGENT": "1"}, {"agent": "ai_agent"})
+        await self._test_handshake({"AI_AGENT": "true"}, {"agent": "ai_agent"})
+
+    async def test_handshake_18_agent_generic_boolean_whitespace(self):
+        await self._test_handshake({"AI_AGENT": " true "}, {"agent": "ai_agent"})
+
+    async def test_handshake_19_agent_generic_normalized(self):
+        await self._test_handshake(
+            {"AI_AGENT": " Claude-Code_2-1-238_Agent "}, {"agent": "claude-code_2-1-238_agent"}
+        )
+
+    async def test_handshake_20_agent_generic_truncated(self):
+        await self._test_handshake({"AI_AGENT": "a" * 100}, {"agent": "a" * _MAX_AGENT_SIZE})
+
+    async def test_handshake_21_agent_generic_truncated_on_boundary(self):
+        # The byte limit falls inside the two-byte "é". 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_22_agent_generic_malformed(self):
+        await self._test_handshake({"AI_AGENT": "custom-\udcff"}, {"agent": "custom-"})
+
+    async def test_handshake_23_agent_unset_empty(self):
+        await self._test_handshake({"AI_AGENT": ""}, None)
+
+    async def test_handshake_24_agent_unset_whitespace(self):
+        await self._test_handshake({"AI_AGENT": "   "}, None)
+
+    async def test_handshake_25_agent_unset_known_whitespace(self):
+        await self._test_handshake({"CLAUDECODE": "   "}, None)
+
+    async def test_handshake_26_agent_with_provider(self):
+        await self._test_handshake(
+            {
+                "AWS_EXECUTION_ENV": "AWS_Lambda_java8",
+                "AWS_REGION": "us-east-2",
+                "CLAUDECODE": "1",
+            },
+            {"name": "aws.lambda", "region": "us-east-2", "agent": "claude_code"},
+        )
+
     def test_dict_hints(self):
         self.db.coll.find(hint={"x": 1})