From 67a8c1fb8f92149b15cf124aa8175416e4c81899 Mon Sep 17 00:00:00 2001 From: Jeffrey 'Alex' Clark Date: Mon, 20 Jul 2026 11:12:58 -0400 Subject: [PATCH 1/6] PYTHON-5929 Add coding agent env var to handshake metadata Detect coding agent environment variables (AI_AGENT, AGENT, CLAUDECODE, CURSOR_AGENT, GEMINI_CLI, CODEX_SANDBOX, AUGMENT_AGENT, OPENCODE_CLIENT) and report them in the client.env.agent handshake field. See DRIVERS-3529. --- pymongo/pool_options.py | 40 ++++++++++++++++++++++++++++---- test/asynchronous/test_client.py | 19 +++++++++++++++ test/test_client.py | 19 +++++++++++++++ 3 files changed, 74 insertions(+), 4 deletions(-) diff --git a/pymongo/pool_options.py b/pymongo/pool_options.py index 8b26b4baf2..fac6c1019c 100644 --- a/pymongo/pool_options.py +++ b/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. +_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,11 @@ 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 if len(bson.encode(metadata)) <= _MAX_METADATA_SIZE: return # 2. Omit fields from os except os.type. diff --git a/test/asynchronous/test_client.py b/test/asynchronous/test_client.py index f26f3fb85b..61b58ce4fd 100644 --- a/test/asynchronous/test_client.py +++ b/test/asynchronous/test_client.py @@ -2204,6 +2204,25 @@ 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_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"}, + ) + def test_dict_hints(self): self.db.t.find(hint={"x": 1}) diff --git a/test/test_client.py b/test/test_client.py index c801d3a178..f113c6a7b9 100644 --- a/test/test_client.py +++ b/test/test_client.py @@ -2161,6 +2161,25 @@ 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_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_dict_hints(self): self.db.t.find(hint={"x": 1}) From 45b85d5badcd6204d2dc5c3924ca65e959ad0090 Mon Sep 17 00:00:00 2001 From: Jeffrey 'Alex' Clark Date: Mon, 20 Jul 2026 11:31:33 -0400 Subject: [PATCH 2/6] PYTHON-5929 Reference PYTHON-5929 in agent env var comment --- pymongo/pool_options.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymongo/pool_options.py b/pymongo/pool_options.py index fac6c1019c..d504cc1e1a 100644 --- a/pymongo/pool_options.py +++ b/pymongo/pool_options.py @@ -151,7 +151,7 @@ def _is_faas() -> bool: # 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. +# See DRIVERS-3529 and PYTHON-5929. _AGENT_ENV_VARS = [ ("CLAUDECODE", "CLAUDECODE"), ("CURSOR_AGENT", "CURSOR"), From 22af204745cfc532c3baa37e081372b2027d08c2 Mon Sep 17 00:00:00 2001 From: Jeffrey 'Alex' Clark Date: Mon, 20 Jul 2026 12:28:25 -0400 Subject: [PATCH 3/6] PYTHON-5929 Drop env.agent before env.name during metadata truncation --- pymongo/pool_options.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/pymongo/pool_options.py b/pymongo/pool_options.py index d504cc1e1a..216746f840 100644 --- a/pymongo/pool_options.py +++ b/pymongo/pool_options.py @@ -249,6 +249,16 @@ def _truncate_metadata(metadata: MutableMapping[str, Any]) -> None: metadata["os"] = {"type": os_type} if len(bson.encode(metadata)) <= _MAX_METADATA_SIZE: return + # 2b. Drop env.agent (which may hold an arbitrarily large AI_AGENT/AGENT + # value) before sacrificing env.name by dropping the env document entirely. + 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 # 3. Omit the env document entirely. metadata.pop("env", None) encoded_size = len(bson.encode(metadata)) From faa732af1a80ab38d3488dd914185343ec66e68c Mon Sep 17 00:00:00 2001 From: Jeffrey 'Alex' Clark Date: Mon, 20 Jul 2026 12:28:52 -0400 Subject: [PATCH 4/6] PYTHON-5929 Add regression test for agent metadata truncation --- test/asynchronous/test_client.py | 7 +++++++ test/test_client.py | 7 +++++++ 2 files changed, 14 insertions(+) diff --git a/test/asynchronous/test_client.py b/test/asynchronous/test_client.py index 61b58ce4fd..ac68e3c290 100644 --- a/test/asynchronous/test_client.py +++ b/test/asynchronous/test_client.py @@ -2223,6 +2223,13 @@ async def test_handshake_12_agent_with_provider(self): {"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}) diff --git a/test/test_client.py b/test/test_client.py index f113c6a7b9..f1c9c2d0f8 100644 --- a/test/test_client.py +++ b/test/test_client.py @@ -2180,6 +2180,13 @@ def test_handshake_12_agent_with_provider(self): {"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}) From ce2cbd7da394c1f1369321a4eb4a17700a9b6c39 Mon Sep 17 00:00:00 2001 From: Jeffrey 'Alex' Clark Date: Wed, 22 Jul 2026 14:04:01 -0400 Subject: [PATCH 5/6] PYTHON-5929 Address Copilot review: drop empty env and isolate agent test env --- pymongo/pool_options.py | 19 +++++++++++-------- test/asynchronous/test_client.py | 14 ++++++++++++-- test/test_client.py | 14 ++++++++++++-- 3 files changed, 35 insertions(+), 12 deletions(-) diff --git a/pymongo/pool_options.py b/pymongo/pool_options.py index 216746f840..37e3cf215d 100644 --- a/pymongo/pool_options.py +++ b/pymongo/pool_options.py @@ -241,16 +241,13 @@ def _truncate_metadata(metadata: MutableMapping[str, Any]) -> None: 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 - # 2. Omit fields from os except os.type. - os_type = metadata.get("os", {}).get("type") - if os_type: - metadata["os"] = {"type": os_type} - if len(bson.encode(metadata)) <= _MAX_METADATA_SIZE: - return - # 2b. Drop env.agent (which may hold an arbitrarily large AI_AGENT/AGENT - # value) before sacrificing env.name by dropping the env document entirely. + # 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: @@ -259,6 +256,12 @@ def _truncate_metadata(metadata: MutableMapping[str, Any]) -> None: metadata.pop("env", None) if len(bson.encode(metadata)) <= _MAX_METADATA_SIZE: return + # 2. Omit fields from os except os.type. + os_type = metadata.get("os", {}).get("type") + if os_type: + metadata["os"] = {"type": os_type} + if len(bson.encode(metadata)) <= _MAX_METADATA_SIZE: + return # 3. Omit the env document entirely. metadata.pop("env", None) encoded_size = len(bson.encode(metadata)) diff --git a/test/asynchronous/test_client.py b/test/asynchronous/test_client.py index ac68e3c290..d6f4743ba6 100644 --- a/test/asynchronous/test_client.py +++ b/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" diff --git a/test/test_client.py b/test/test_client.py index f1c9c2d0f8..103095d6a8 100644 --- a/test/test_client.py +++ b/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" From 62d8a88cffaa2838ddfcabed8254867dad6bfe12 Mon Sep 17 00:00:00 2001 From: Jeffrey 'Alex' Clark Date: Wed, 22 Jul 2026 14:32:08 -0400 Subject: [PATCH 6/6] PYTHON-5929 Assert known-agent env var precedence in handshake test --- test/asynchronous/test_client.py | 7 +++++++ test/test_client.py | 7 +++++++ 2 files changed, 14 insertions(+) diff --git a/test/asynchronous/test_client.py b/test/asynchronous/test_client.py index d6f4743ba6..2b33e80dd9 100644 --- a/test/asynchronous/test_client.py +++ b/test/asynchronous/test_client.py @@ -2220,6 +2220,13 @@ async def test_handshake_10_agent_known(self): 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"}) diff --git a/test/test_client.py b/test/test_client.py index 103095d6a8..c932df582d 100644 --- a/test/test_client.py +++ b/test/test_client.py @@ -2177,6 +2177,13 @@ def test_handshake_10_agent_known(self): 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"})