PYTHON-5929 Add coding agent env var to handshake metadata

pymongo/pool_options.py

@@ -149,6 +149,34 @@ def _is_faas() -> bool:
     return _is_lambda() or _is_azure_func() or _is_gcp_func() or _is_vercel()
 
 
+# Environment variables that indicate a coding agent, checked in order. The
+# first match determines the value of the client.env.agent metadata field.
+# See DRIVERS-3529 and PYTHON-5929.
+_AGENT_ENV_VARS = [
+    ("CLAUDECODE", "CLAUDECODE"),
+    ("CURSOR_AGENT", "CURSOR"),
+    ("GEMINI_CLI", "GEMINI_CLI"),
+    ("CODEX_SANDBOX", "CODEX_SANDBOX"),
+    ("AUGMENT_AGENT", "AUGMENT"),
+    ("OPENCODE_CLIENT", "OPENCODE"),
+]
+
+
+def _metadata_agent() -> Optional[str]:
+    """Detect a coding agent from the environment for client.env.agent.
+
+    A generic AI_AGENT or AGENT environment variable takes precedence and its
+    value is used verbatim. Otherwise the first matching known agent variable
+    determines the value."""
+    agent = os.getenv("AI_AGENT") or os.getenv("AGENT")
+    if agent:
+        return agent
+    for var, name in _AGENT_ENV_VARS:
+        if os.getenv(var):
+            return name
+    return None
+
+
 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)
@@ -165,6 +193,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
@@ -205,10 +236,24 @@ def _truncate_metadata(metadata: MutableMapping[str, Any]) -> None:
     """Perform metadata truncation."""
     if len(bson.encode(metadata)) <= _MAX_METADATA_SIZE:
         return
-    # 1. Omit fields from env except env.name.
-    env_name = metadata.get("env", {}).get("name")
-    if env_name:
-        metadata["env"] = {"name": env_name}
+    # 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:
+        metadata["env"] = trimmed_env
+    else:
+        metadata.pop("env", None)
+    if len(bson.encode(metadata)) <= _MAX_METADATA_SIZE:
+        return
+    # 1b. Drop env.agent (which may hold an arbitrarily large AI_AGENT/AGENT
+    # value) before trimming os and before sacrificing env.name, so the more
+    # valuable os and env.name fields are preserved as long as possible.
+    if "agent" in trimmed_env:
+        del trimmed_env["agent"]
+        if trimmed_env:
+            metadata["env"] = trimmed_env
+        else:
+            metadata.pop("env", None)
     if len(bson.encode(metadata)) <= _MAX_METADATA_SIZE:
         return
     # 2. Omit fields from os except os.type.

test/asynchronous/test_client.py

@@ -87,7 +87,13 @@
     WriteConcernError,
 )
 from pymongo.monitoring import ServerHeartbeatListener, ServerHeartbeatStartedEvent
-from pymongo.pool_options import _MAX_METADATA_SIZE, _METADATA, ENV_VAR_K8S, PoolOptions
+from pymongo.pool_options import (
+    _AGENT_ENV_VARS,
+    _MAX_METADATA_SIZE,
+    _METADATA,
+    ENV_VAR_K8S,
+    PoolOptions,
+)
 from pymongo.read_preferences import ReadPreference
 from pymongo.server_description import ServerDescription
 from pymongo.server_selectors import readable_server_selector, writable_server_selector
@@ -2099,7 +2105,11 @@ 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 any ambient agent-detection vars (e.g. AI_AGENT/AGENT set by the
+        # CI runner) so detection is deterministic and only reflects env_vars.
+        agent_vars = ["AI_AGENT", "AGENT", *(var for var, _ in _AGENT_ENV_VARS)]
+        cleared = {var: "" for var in agent_vars if var not in env_vars}
+        with patch.dict("os.environ", {**cleared, **env_vars}):
             metadata = copy.deepcopy(_METADATA)
             if has_c():
                 metadata["driver"]["name"] = "PyMongo|c|async"
@@ -2204,6 +2214,39 @@ async def test_handshake_09_container_with_provider(self):
             },
         )
 
+    async def test_handshake_10_agent_known(self):
+        # A known coding-agent env var maps to its metadata value.
+        await self._test_handshake({"CLAUDECODE": "1"}, {"agent": "CLAUDECODE"})
+        await self._test_handshake({"CURSOR_AGENT": "1"}, {"agent": "CURSOR"})
+        await self._test_handshake({"OPENCODE_CLIENT": "1"}, {"agent": "OPENCODE"})
+
+    async def test_handshake_10b_agent_known_precedence(self):
+        # When multiple known agent vars are set, the first in _AGENT_ENV_VARS
+        # order wins, regardless of which comes first in the environment dict.
+        first_var, first_name = _AGENT_ENV_VARS[0]
+        last_var, _ = _AGENT_ENV_VARS[-1]
+        await self._test_handshake({last_var: "1", first_var: "1"}, {"agent": first_name})
+
+    async def test_handshake_11_agent_generic(self):
+        # Generic AI_AGENT/AGENT vars are used verbatim and take precedence.
+        await self._test_handshake({"AI_AGENT": "myagent"}, {"agent": "myagent"})
+        await self._test_handshake({"AGENT": "myagent"}, {"agent": "myagent"})
+        await self._test_handshake({"AI_AGENT": "myagent", "CLAUDECODE": "1"}, {"agent": "myagent"})
+
+    async def test_handshake_12_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": "CLAUDECODE"},
+        )
+
+    async def test_handshake_13_agent_too_long(self):
+        # A too-long agent value is dropped during truncation before env.name.
+        await self._test_handshake(
+            {"FUNCTIONS_WORKER_RUNTIME": "python", "AI_AGENT": "a" * 512},
+            {"name": "azure.func"},
+        )
+
     def test_dict_hints(self):
         self.db.t.find(hint={"x": 1})
 

test/test_client.py

@@ -76,7 +76,13 @@
     WriteConcernError,
 )
 from pymongo.monitoring import ServerHeartbeatListener, ServerHeartbeatStartedEvent
-from pymongo.pool_options import _MAX_METADATA_SIZE, _METADATA, ENV_VAR_K8S, PoolOptions
+from pymongo.pool_options import (
+    _AGENT_ENV_VARS,
+    _MAX_METADATA_SIZE,
+    _METADATA,
+    ENV_VAR_K8S,
+    PoolOptions,
+)
 from pymongo.read_preferences import ReadPreference
 from pymongo.server_description import ServerDescription
 from pymongo.server_selectors import readable_server_selector, writable_server_selector
@@ -2056,7 +2062,11 @@ 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 any ambient agent-detection vars (e.g. AI_AGENT/AGENT set by the
+        # CI runner) so detection is deterministic and only reflects env_vars.
+        agent_vars = ["AI_AGENT", "AGENT", *(var for var, _ in _AGENT_ENV_VARS)]
+        cleared = {var: "" for var in agent_vars if var not in env_vars}
+        with patch.dict("os.environ", {**cleared, **env_vars}):
             metadata = copy.deepcopy(_METADATA)
             if has_c():
                 metadata["driver"]["name"] = "PyMongo|c"
@@ -2161,6 +2171,39 @@ def test_handshake_09_container_with_provider(self):
             },
         )
 
+    def test_handshake_10_agent_known(self):
+        # A known coding-agent env var maps to its metadata value.
+        self._test_handshake({"CLAUDECODE": "1"}, {"agent": "CLAUDECODE"})
+        self._test_handshake({"CURSOR_AGENT": "1"}, {"agent": "CURSOR"})
+        self._test_handshake({"OPENCODE_CLIENT": "1"}, {"agent": "OPENCODE"})
+
+    def test_handshake_10b_agent_known_precedence(self):
+        # When multiple known agent vars are set, the first in _AGENT_ENV_VARS
+        # order wins, regardless of which comes first in the environment dict.
+        first_var, first_name = _AGENT_ENV_VARS[0]
+        last_var, _ = _AGENT_ENV_VARS[-1]
+        self._test_handshake({last_var: "1", first_var: "1"}, {"agent": first_name})
+
+    def test_handshake_11_agent_generic(self):
+        # Generic AI_AGENT/AGENT vars are used verbatim and take precedence.
+        self._test_handshake({"AI_AGENT": "myagent"}, {"agent": "myagent"})
+        self._test_handshake({"AGENT": "myagent"}, {"agent": "myagent"})
+        self._test_handshake({"AI_AGENT": "myagent", "CLAUDECODE": "1"}, {"agent": "myagent"})
+
+    def test_handshake_12_agent_with_provider(self):
+        # agent is reported alongside a FaaS provider.
+        self._test_handshake(
+            {"FUNCTIONS_WORKER_RUNTIME": "python", "CLAUDECODE": "1"},
+            {"name": "azure.func", "agent": "CLAUDECODE"},
+        )
+
+    def test_handshake_13_agent_too_long(self):
+        # A too-long agent value is dropped during truncation before env.name.
+        self._test_handshake(
+            {"FUNCTIONS_WORKER_RUNTIME": "python", "AI_AGENT": "a" * 512},
+            {"name": "azure.func"},
+        )
+
     def test_dict_hints(self):
         self.db.t.find(hint={"x": 1})