(WIP) feat(NODE-7875): environment variables in the handshake metadata

src/cmap/handshake/client_metadata.ts

@@ -1,9 +1,9 @@
 import * as process from 'process';
 
-import { BSON, ByteUtils, type Document, Int32, NumberUtils } from '../../bson';
+import { BSON, type Document, Int32, NumberUtils } from '../../bson';
 import { MongoInvalidArgumentError } from '../../error';
 import type { DriverInfo, MongoOptions } from '../../mongo_client';
-import { fileIsAccessible } from '../../utils';
+import { DriverStringUtils, fileIsAccessible } from '../../utils';
 
 // eslint-disable-next-line @typescript-eslint/no-require-imports
 const NODE_DRIVER_VERSION = require('../../../package.json').version;
@@ -44,16 +44,17 @@ export interface ClientMetadata {
   application?: {
     name: string;
   };
-  /** FaaS environment information */
   env?: {
-    name?: 'aws.lambda' | 'gcp.func' | 'azure.func' | 'vercel';
-    timeout_sec?: Int32;
-    memory_mb?: Int32;
-    region?: string;
+    agent?: string;
     container?: {
       runtime?: string;
       orchestrator?: string;
     };
+    /** FaaS environment information */
+    name?: 'aws.lambda' | 'gcp.func' | 'azure.func' | 'vercel';
+    timeout_sec?: Int32;
+    memory_mb?: Int32;
+    region?: string;
   };
 }
 
@@ -74,11 +75,18 @@ export class LimitedSizeDocument {
     // subtracting the document size int32 and the null terminator.
     const newElementSize = BSON.serialize(new Map().set(key, value)).byteLength - 5;
 
-    if (newElementSize + this.documentSize > this.maxSize) {
+    let baseDocumentSize = this.documentSize;
+    // If the incoming key is a replace op, don't double-count the size
+    if (this.document.has(key)) {
+      // new document size = current size - existing value, prepping for replacement
+      baseDocumentSize -= BSON.serialize(new Map().set(key, this.document.get(key))).byteLength - 5;
+    }
+
+    if (newElementSize + baseDocumentSize > this.maxSize) {
       return false;
     }
 
-    this.documentSize += newElementSize;
+    this.documentSize = baseDocumentSize + newElementSize;
 
     this.document.set(key, value);
 
@@ -95,12 +103,15 @@ export class LimitedSizeDocument {
   }
 }
 
+// Env var value length after truncating bytes
+const APP_NAME_TRUNCATION_LIMIT_BYTES = 128;
+
 type MakeClientMetadataOptions = Pick<MongoOptions, 'appName' | 'runtime'>;
 
 /**
  * From the specs:
  * Implementors SHOULD cumulatively update fields in the following order until the document is under the size limit:
- * 1. Omit fields from `env` except `env.name`.
+ * 1. Omit fields from `env` except `env.name` & `env.agent`.
  * 2. Omit fields from `os` except `os.type`.
  * 3. Omit the `env` document entirely.
  * 4. Truncate `platform`. -- special we do not truncate this field
@@ -114,10 +125,7 @@ export async function makeClientMetadata(
 
   // Add app name first, it must be sent
   if (appName.length > 0) {
-    const name =
-      ByteUtils.utf8ByteLength(appName) <= 128
-        ? appName
-        : ByteUtils.toUTF8(ByteUtils.fromUTF8(appName), 0, 128, false);
+    const name = DriverStringUtils.truncateStringBytes(appName, APP_NAME_TRUNCATION_LIMIT_BYTES);
     metadataDocument.ifItFitsItSits('application', { name });
   }
 
@@ -171,76 +179,138 @@ export async function makeClientMetadata(
     }
   }
 
+  // TODO(NODE-7851): Env has an order of precedence for data truncation, and order matters.
+  // We append in the order of delete preference. 'name' is appended at the end of
+  // faasEnv, and is preference-agnostic with 'agent' for now, since ensuring
+  // truncation order is to be tackled in NODE-7851.
+  const containerMetadata = await getContainerMetadata();
   const faasEnv = getFAASEnv();
-  if (faasEnv != null) {
-    if (!metadataDocument.ifItFitsItSits('env', faasEnv)) {
-      for (const key of faasEnv.keys()) {
-        faasEnv.delete(key);
-        if (faasEnv.size === 0) break;
-        if (metadataDocument.ifItFitsItSits('env', faasEnv)) break;
-      }
+  const agentEnv = getAgentEnv();
+
+  const fullEnv = new Map<string, unknown>();
+  for (const [k, v] of faasEnv) fullEnv.set(k, v);
+  if (agentEnv.length > 0) fullEnv.set('agent', agentEnv);
+  if (fullEnv.size > 0 && !metadataDocument.ifItFitsItSits('env', fullEnv)) {
+    for (const key of fullEnv.keys()) {
+      fullEnv.delete(key);
+      if (fullEnv.size === 0) break;
+      if (metadataDocument.ifItFitsItSits('env', fullEnv)) break;
     }
   }
-  return await addContainerMetadata(metadataDocument.toObject() as ClientMetadata);
+  if (containerMetadata.size > 0) {
+    const newEnv = { ...Object.fromEntries(fullEnv), container: containerMetadata };
+    metadataDocument.ifItFitsItSits('env', newEnv);
+  }
+
+  return metadataDocument.toObject() as ClientMetadata;
 }
 
-let dockerPromise: Promise<boolean>;
-type ContainerMetadata = NonNullable<NonNullable<ClientMetadata['env']>['container']>;
+let dockerPromise: Promise<boolean> | undefined;
+
 /** @internal */
-async function getContainerMetadata(): Promise<ContainerMetadata> {
+export function resetDockerPromise(value?: Promise<boolean>) {
+  dockerPromise = value;
+}
+
+/** @internal */
+async function getContainerMetadata(): Promise<Map<string, string>> {
   dockerPromise ??= fileIsAccessible('/.dockerenv');
   const isDocker = await dockerPromise;
 
   const { KUBERNETES_SERVICE_HOST = '' } = process.env;
   const isKubernetes = KUBERNETES_SERVICE_HOST.length > 0 ? true : false;
 
-  const containerMetadata: ContainerMetadata = {};
+  const containerMetadata = new Map<string, string>();
 
-  if (isDocker) containerMetadata.runtime = 'docker';
-  if (isKubernetes) containerMetadata.orchestrator = 'kubernetes';
+  if (isDocker) containerMetadata.set('runtime', 'docker');
+  if (isKubernetes) containerMetadata.set('orchestrator', 'kubernetes');
 
   return containerMetadata;
 }
 
 /**
  * @internal
- * Re-add each metadata value.
- * Attempt to add new env container metadata, but keep old data if it does not fit.
+ * Environment variables that indicate the driver is being used by an AI agent, in the order the
+ * spec requires them to be evaluated. [0] is the environment variable, [1] is the value to set
+ * when that environment variable is encountered. If [1] is null, use the environment variable value.
+ *
+ * From the spec:
+ * client.env.agent is a single string. Its value is determined by the environment variables below.
+ * Drivers MUST evaluate the list in order. The first populated variable determines the value, and
+ * subsequent entries MUST NOT be considered.
  */
-async function addContainerMetadata(originalMetadata: ClientMetadata): Promise<ClientMetadata> {
-  const containerMetadata = await getContainerMetadata();
-  if (Object.keys(containerMetadata).length === 0) return originalMetadata;
+export const AGENT_ENV_VARIABLES: ReadonlyArray<readonly [string, string | null]> = [
+  ['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'],
+  ['AI_AGENT', null]
+];
+
+// Env var value length after first 2 forms of normalization
+const AGENT_ENV_LIMIT_BYTES = 64;
 
-  const extendedMetadata = new LimitedSizeDocument(512);
-
-  const extendedEnvMetadata: NonNullable<ClientMetadata['env']> = {
-    ...originalMetadata?.env,
-    container: containerMetadata
-  };
+/**
+ * @internal
+ * Values used in AI_AGENT to indicate an agent is being used but are non-identifying
+ */
+export const AGENT_ENV_UNIDENTIFYING = new Set(['1', 'true']);
 
-  for (const [key, val] of Object.entries(originalMetadata)) {
-    if (key !== 'env') {
-      extendedMetadata.ifItFitsItSits(key, val);
-    } else {
-      if (!extendedMetadata.ifItFitsItSits('env', extendedEnvMetadata)) {
-        // add in old data if newer / extended metadata does not fit
-        extendedMetadata.ifItFitsItSits('env', val);
+/**
+ * @internal
+ * Resolves `env.agent` from the environment, or an empty string when no agent variable is
+ * populated. Returns the value of the first populated variable in `AGENT_ENV_VARIABLES`.
+ */
+export function getAgentEnv(): string {
+  for (const [key, literal] of AGENT_ENV_VARIABLES) {
+    // A variable is only populated if it is present with a non-empty normalized value, so an empty or
+    // whitespace-only value never selects an entry, even one with a fixed table value.
+    let envValue = (process.env[key] ?? '').trim().toLowerCase();
+
+    if (envValue.length > 0) {
+      envValue = DriverStringUtils.truncateStringBytes(envValue, AGENT_ENV_LIMIT_BYTES);
+      if (!literal) {
+        return AGENT_ENV_UNIDENTIFYING.has(envValue) ? 'ai_agent' : envValue;
       }
+      return literal;
     }
   }
 
-  if (!('env' in originalMetadata)) {
-    extendedMetadata.ifItFitsItSits('env', extendedEnvMetadata);
-  }
-
-  return extendedMetadata.toObject() as ClientMetadata;
+  return '';
 }
 
+/**
+ * @internal
+ * Environment variables relevant to FaaS runtimes, governed by the spec:
+ * https://github.com/mongodb/specifications/blob/9cfe388c7d2ce1e02b24b53606e97ce7b73ebb7d/source/mongodb-handshake/handshake.md#faas
+ */
+export const FAAS_ENV_VARIABLES = [
+  'AWS_EXECUTION_ENV',
+  'AWS_LAMBDA_RUNTIME_API',
+  'FUNCTIONS_WORKER_RUNTIME',
+  'K_SERVICE',
+  'FUNCTION_NAME',
+  'VERCEL',
+  'AWS_LAMBDA_FUNCTION_MEMORY_SIZE',
+  'AWS_REGION',
+  'FUNCTION_MEMORY_MB',
+  'FUNCTION_REGION',
+  'FUNCTION_TIMEOUT_SEC',
+  'VERCEL_REGION'
+] as const;
+
 /**
  * Collects FaaS metadata.
  * - `name` MUST be the last key in the Map returned.
  */
-export function getFAASEnv(): Map<string, string | Int32> | null {
+export function getFAASEnv(): Map<string, string | Int32> {
   const {
     AWS_EXECUTION_ENV = '',
     AWS_LAMBDA_RUNTIME_API = '',
@@ -313,7 +383,7 @@ export function getFAASEnv(): Map<string, string | Int32> | null {
     return faasEnv;
   }
 
-  return null;
+  return faasEnv;
 }
 
 /**

src/sdam/monitor.ts

@@ -120,7 +120,7 @@ export class Monitor extends TypedEventEmitter<MonitorEvents> {
       minHeartbeatFrequencyMS: options.minHeartbeatFrequencyMS ?? 500,
       serverMonitoringMode: options.serverMonitoringMode
     });
-    this.isRunningInFaasEnv = getFAASEnv() != null;
+    this.isRunningInFaasEnv = getFAASEnv().size > 0;
     this.mongoLogger = this.server.topology.client?.mongoLogger;
     this.rttSampler = new RTTSampler(10);
 

src/utils.ts

@@ -3,6 +3,7 @@ import { type EventEmitter } from 'events';
 import { promises as fs } from 'fs';
 import * as http from 'http';
 import * as process from 'process';
+import { StringDecoder } from 'string_decoder';
 import { clearTimeout, setTimeout } from 'timers';
 
 import {
@@ -1143,6 +1144,25 @@ export function parseUnsignedInteger(value: unknown): number | null {
   return parsedInt != null && parsedInt >= 0 ? parsedInt : null;
 }
 
+/** @internal */
+export const DriverStringUtils = {
+  /**
+   * Truncates a string to a specified byte length. Multi-byte characters which are split by the byte truncation
+   * are considered, ensuring incomplete multi-byte characters are omitted from the result entirely.
+   * @throws MongoRuntimeError on a malformed limit
+   */
+  truncateStringBytes(str: string, limit: number): string {
+    if (!Number.isInteger(limit) || limit < 0) {
+      throw new MongoRuntimeError('provided limit must be a non-negative integer');
+    }
+    const buf = new TextEncoder().encode(str);
+    if (buf.length <= limit) return str;
+    const decoder = new StringDecoder('utf8');
+    // incomplete multi-byte chars will not be returned through decoder.write
+    return decoder.write(buf.subarray(0, limit));
+  }
+};
+
 /**
  * This function throws a MongoAPIError in the event that either of the following is true:
  * * If the provided address domain does not match the provided parent domain