Add telemetry correlation headers to dev tunnel service requests

cli/CONTRIBUTING.md

@@ -31,3 +31,18 @@ an interactive prompt.
 
 Without `--tunnel`, `code agent` and `code agent host` still start or reuse a local
 agent host supervisor.
+
+## Tunnel request correlation
+
+Management requests and relay WebSocket handshakes carry the
+`X-Tunnels-VSCode-Session-Id`, `X-Tunnels-VSCode-Client-Operation-Id`, and
+`X-Tunnels-VSCode-Client-Request-Id` headers. Each request has a fresh request ID.
+VS Code passes its telemetry session and operation IDs through
+`VSCODE_TUNNEL_SESSION_ID` and `VSCODE_TUNNEL_OPERATION_ID` when launching the CLI.
+These IDs are preserved verbatim; telemetry session IDs are not necessarily UUIDs.
+Standalone commands use a process-local session ID and a new operation ID for each
+management client. Correlation IDs are sent in service headers, not written to CLI logs.
+Local CLI status, login, and service-management invocations retain correlation IDs but
+do not emit host-operation telemetry.
+Attaching to an already-running tunnel does not replace that process's correlation IDs;
+its service requests retain the IDs of the original host process.

cli/Cargo.toml

@@ -33,7 +33,7 @@ jiff = { version = "0.2", default-features = false, features = ["std", "serde"]
 http = "1"
 gethostname = "0.4.3"
 libc = "0.2.144"
-tunnels = { git = "https://github.com/microsoft/dev-tunnels", rev = "64048c1409ff56cb958b879de7ea069ec71edc8b", default-features = false, features = ["connections"] }
+tunnels = { git = "https://github.com/microsoft/dev-tunnels", rev = "b73b0eb774b86a5b519c300f2f41c9072a66f663", default-features = false, features = ["connections"] }
 keyring = { version = "2.0.3", default-features = false, features = ["linux-secret-service-rt-tokio-crypto-openssl", "platform-windows", "platform-macos", "linux-keyutils"] }
 dialoguer = "0.10.4"
 hyper = { version = "1", features = ["server", "http1", "client"] }

cli/src/tunnels/dev_tunnels.rs

@@ -14,6 +14,7 @@ use crate::util::errors::{
 use crate::util::input::prompt_placeholder;
 use futures::future::BoxFuture;
 use futures::{FutureExt, TryFutureExt};
+use http::header::{HeaderName, HeaderValue};
 use http::StatusCode;
 use rand::prelude::IteratorRandom;
 use regex::Regex;
@@ -30,12 +31,73 @@ use tunnels::contracts::{
 	TUNNEL_ACCESS_SCOPES_CONNECT, TUNNEL_PROTOCOL_AUTO,
 };
 use tunnels::management::{
-	new_tunnel_management, HttpError, TunnelLocator, TunnelManagementClient, TunnelRequestOptions,
-	NO_REQUEST_OPTIONS,
+	new_tunnel_management, HttpError, TunnelClientBuilder, TunnelLocator, TunnelManagementClient,
+	TunnelRequestOptions, NO_REQUEST_OPTIONS,
 };
 
 static TUNNEL_COUNT_LIMIT_NAME: &str = "TunnelsPerUserPerLocation";
 static TUNNEL_PORT_PROTOCOL_CONFLICT_DETAIL: &str = "The tunnel port protocol cannot be changed.";
+static TUNNEL_SESSION_ID: LazyLock<uuid::Uuid> = LazyLock::new(uuid::Uuid::new_v4);
+
+fn tunnel_correlation_id(
+	log: &log::Logger,
+	name: &str,
+	value: Result<String, std::env::VarError>,
+	fallback: uuid::Uuid,
+) -> HeaderValue {
+	let fallback = HeaderValue::from_str(&fallback.to_string()).unwrap();
+	match value {
+		Ok(value) if !value.is_empty() => match HeaderValue::from_str(&value) {
+			Ok(id) => id,
+			Err(error) => {
+				warning!(log, "Invalid {}: {}", name, error);
+				fallback
+			}
+		},
+		Ok(_) => {
+			warning!(log, "Invalid {}: empty correlation ID", name);
+			fallback
+		}
+		Err(std::env::VarError::NotPresent) => fallback,
+		Err(error) => {
+			warning!(log, "Invalid {}: {}", name, error);
+			fallback
+		}
+	}
+}
+
+fn tunnel_management_client(log: &log::Logger) -> TunnelClientBuilder {
+	let session_id = tunnel_correlation_id(
+		log,
+		"VSCODE_TUNNEL_SESSION_ID",
+		std::env::var("VSCODE_TUNNEL_SESSION_ID"),
+		*TUNNEL_SESSION_ID,
+	);
+	let operation_id = tunnel_correlation_id(
+		log,
+		"VSCODE_TUNNEL_OPERATION_ID",
+		std::env::var("VSCODE_TUNNEL_OPERATION_ID"),
+		uuid::Uuid::new_v4(),
+	);
+	let mut client = new_tunnel_management(&TUNNEL_SERVICE_USER_AGENT);
+	client.additional_headers(vec![
+		(
+			HeaderName::from_static("x-tunnels-vscode-session-id"),
+			session_id,
+		),
+		(
+			HeaderName::from_static("x-tunnels-vscode-client-operation-id"),
+			operation_id,
+		),
+	]);
+	client.additional_headers_provider(|| {
+		vec![(
+			HeaderName::from_static("x-tunnels-vscode-client-request-id"),
+			HeaderValue::from_str(&uuid::Uuid::new_v4().to_string()).unwrap(),
+		)]
+	});
+	client
+}
 
 #[allow(dead_code)]
 mod tunnel_flags {
@@ -329,7 +391,7 @@ impl DevTunnels {
 		auth: auth::Auth,
 		paths: &LauncherPaths,
 	) -> DevTunnels {
-		let mut client = new_tunnel_management(&TUNNEL_SERVICE_USER_AGENT);
+		let mut client = tunnel_management_client(log);
 		client.authorization_provider(auth.clone());
 
 		DevTunnels {
@@ -347,7 +409,7 @@ impl DevTunnels {
 		auth: auth::Auth,
 		paths: &LauncherPaths,
 	) -> DevTunnels {
-		let mut client = new_tunnel_management(&TUNNEL_SERVICE_USER_AGENT);
+		let mut client = tunnel_management_client(log);
 		client.authorization_provider(auth.clone());
 
 		DevTunnels {
@@ -1328,6 +1390,59 @@ fn tunnel_has_host_connection(tunnel: &Tunnel) -> bool {
 mod test {
 	use super::*;
 
+	#[test]
+	fn test_tunnel_correlation_preserves_telemetry_session_id() {
+		let session_id = "7c86e247-3438-4f26-959d-061b3299f4de1791567000000";
+		assert_eq!(
+			tunnel_correlation_id(
+				&log::Logger::test(),
+				"VSCODE_TUNNEL_SESSION_ID",
+				Ok(session_id.to_string()),
+				uuid::Uuid::new_v4(),
+			),
+			HeaderValue::from_static(session_id),
+		);
+	}
+
+	#[test]
+	fn test_tunnel_correlation_rejects_invalid_header_values() {
+		let fallback = uuid::Uuid::new_v4();
+		for value in ["", "session\r\ninjected-header: value"] {
+			assert_eq!(
+				tunnel_correlation_id(
+					&log::Logger::test(),
+					"VSCODE_TUNNEL_SESSION_ID",
+					Ok(value.to_string()),
+					fallback,
+				),
+				HeaderValue::from_str(&fallback.to_string()).unwrap(),
+			);
+		}
+	}
+
+	#[test]
+	fn test_tunnel_request_correlation() {
+		let client: TunnelManagementClient = tunnel_management_client(&log::Logger::test()).into();
+		let first = client.request_headers();
+		let second = client.request_headers();
+		assert_eq!(
+			first
+				.iter()
+				.map(|(name, _)| name.as_str())
+				.collect::<Vec<_>>(),
+			vec![
+				"x-tunnels-vscode-session-id",
+				"x-tunnels-vscode-client-operation-id",
+				"x-tunnels-vscode-client-request-id",
+			]
+		);
+		for (_, value) in &first {
+			assert!(uuid::Uuid::parse_str(value.to_str().unwrap()).is_ok());
+		}
+		assert_eq!(&first[..2], &second[..2]);
+		assert_ne!(first[2].1, second[2].1);
+	}
+
 	#[test]
 	fn test_clean_hostname_for_tunnel() {
 		assert_eq!(

package.json

@@ -106,9 +106,9 @@
     "@hpke/core": "^1.9.0",
     "@microsoft/1ds-core-js": "^3.2.13",
     "@microsoft/1ds-post-js": "^3.2.13",
-    "@microsoft/dev-tunnels-connections": "^1.3.41",
-    "@microsoft/dev-tunnels-contracts": "^1.3.41",
-    "@microsoft/dev-tunnels-management": "^1.3.41",
+    "@microsoft/dev-tunnels-connections": "^1.3.61",
+    "@microsoft/dev-tunnels-contracts": "^1.3.61",
+    "@microsoft/dev-tunnels-management": "^1.3.61",
     "@microsoft/dev-tunnels-ssh": "^3.12.22",
     "@microsoft/dev-tunnels-ssh-tcp": "^3.12.22",
     "@parcel/watcher": "^2.5.6",

src/vs/platform/agentHost/node/tunnelAgentHostService.ts

@@ -9,6 +9,8 @@ import type WebSocket from 'ws';
 import { Emitter, Event } from '../../../base/common/event.js';
 import { Disposable, toDisposable } from '../../../base/common/lifecycle.js';
 import { ILogService } from '../../log/common/log.js';
+import { ITelemetryService } from '../../telemetry/common/telemetry.js';
+import { createTunnelServiceCorrelation, tunnelServiceHeaders } from '../../remoteTunnel/common/tunnelServiceHeaders.js';
 import {
 	PendingGatewaySelection,
 	TunnelAgentHostConnector,
@@ -185,10 +187,11 @@ export class TunnelAgentHostMainService extends Disposable implements ITunnelAge
 
 	constructor(
 		@ILogService private readonly _logService: ILogService,
+		@ITelemetryService private readonly _telemetryService: ITelemetryService,
 	) {
 		super();
 		this._connector = this._register(new TunnelAgentHostConnector(
-			new NodeTunnelRelayClientFactory((token, authProvider) => this._createManagementClient(token, authProvider)),
+			new NodeTunnelRelayClientFactory((token, authProvider) => this._createManagementClient(token, authProvider, 'connect')),
 			new NodeTunnelSocketFactory(),
 			this._logService,
 		));
@@ -198,7 +201,7 @@ export class TunnelAgentHostMainService extends Disposable implements ITunnelAge
 	}
 
 	async listTunnels(token: string, authProvider: 'github' | 'microsoft', additionalTunnelNames?: string[]): Promise<ITunnelInfo[]> {
-		const client = await this._createManagementClient(token, authProvider);
+		const client = await this._createManagementClient(token, authProvider, 'list');
 		const results: ITunnelInfo[] = [];
 		const seen = new Set<string>();
 
@@ -248,7 +251,7 @@ export class TunnelAgentHostMainService extends Disposable implements ITunnelAge
 	}
 
 	async deleteTunnel(token: string, authProvider: 'github' | 'microsoft', tunnelId: string, clusterId: string): Promise<void> {
-		const client = await this._createManagementClient(token, authProvider);
+		const client = await this._createManagementClient(token, authProvider, 'delete');
 		this._logService.info(`${LOG_PREFIX} Deleting tunnel ${tunnelId} in cluster ${clusterId}...`);
 		await client.deleteTunnel({ tunnelId, clusterId });
 		this._connector.closeTunnelConnections(tunnelId, 'deleting');
@@ -279,14 +282,16 @@ export class TunnelAgentHostMainService extends Disposable implements ITunnelAge
 		return this._connector.disconnect(connectionId);
 	}
 
-	private async _createManagementClient(token: string, authProvider: 'github' | 'microsoft'): Promise<TunnelManagementHttpClient> {
+	private async _createManagementClient(token: string, authProvider: 'github' | 'microsoft', operation: 'list' | 'connect' | 'delete'): Promise<TunnelManagementHttpClient> {
 		const management = await import('@microsoft/dev-tunnels-management');
 		const authHeader = authProvider === 'github' ? `github ${token}` : `Bearer ${token}`;
-		return new management.TunnelManagementHttpClient(
+		const client = new management.TunnelManagementHttpClient(
 			'vscode-sessions',
 			management.ManagementApiVersions.Version20230927preview,
 			async () => authHeader,
 		);
+		client.additionalRequestHeaders = tunnelServiceHeaders(createTunnelServiceCorrelation(this._telemetryService, operation));
+		return client;
 	}
 }