refactor(telemetry): simplify startup telemetry instrumentation

- Bundle activation tracer methods (setAuthState, traceDeploymentInit) into a span-scoped ActivationTracer interface; drop instance-state from ActivationTelemetry
- Mirror the same shape for remote.setup: RemoteSetupTracer exposes phase() and is owned by RemoteSetupTelemetry; remote.ts no longer touches Span directly
- Build RemoteSetupContext after auth retrieval (no empty-default mutation); group setup args under args; rename baseUrlRaw -> baseUrl
- Reuse AuthorityParts from util instead of synthesizing a local alias
- Drop redundant post-download recordDownloadedBytes and the BinaryDownloadResult/DownloadResult types; per-chunk progress already records the final value
- Adopt createTestTelemetryService helper in commandManager.test.ts

Author: EhabY

src/core/cliManager.ts

@@ -36,16 +36,6 @@ type ResolvedBinary =
 
 type CliDownloadReason = "missing" | "version_mismatch";
 
-interface BinaryDownloadResult {
-	binPath: string;
-	downloadedBytes?: number;
-}
-
-interface DownloadResult {
-	status: number;
-	downloadedBytes: number;
-}
-
 export class CliManager {
 	private readonly binaryLock: BinaryLock;
 
@@ -243,22 +233,18 @@ export class CliManager {
 			return await this.telemetry.trace(
 				"cli.download",
 				async (span) => {
-					const recordDownloadedBytes = (downloadedBytes: number): void => {
-						if (downloadedBytes > 0) {
-							span.setMeasurement("downloadedBytes", downloadedBytes);
-						}
-					};
-					const downloadResult = await this.performBinaryDownload(
+					const downloadedBinPath = await this.performBinaryDownload(
 						restClient,
 						latestVersion,
 						downloadBinPath,
 						progressLogPath,
-						recordDownloadedBytes,
+						(downloadedBytes) => {
+							if (downloadedBytes > 0) {
+								span.setMeasurement("downloadedBytes", downloadedBytes);
+							}
+						},
 					);
-					if (downloadResult.downloadedBytes !== undefined) {
-						recordDownloadedBytes(downloadResult.downloadedBytes);
-					}
-					return this.renameToFinalPath(resolved, downloadResult.binPath);
+					return this.renameToFinalPath(resolved, downloadedBinPath);
 				},
 				{ reason: downloadReason },
 			);
@@ -445,7 +431,7 @@ export class CliManager {
 		binPath: string,
 		progressLogPath: string,
 		recordDownloadedBytes: (downloadedBytes: number) => void,
-	): Promise<BinaryDownloadResult> {
+	): Promise<string> {
 		const cfg = vscode.workspace.getConfiguration("coder");
 		const tempFile = tempFilePath(binPath, "temp");
 
@@ -490,7 +476,7 @@ export class CliManager {
 			};
 
 			const client = restClient.getAxiosInstance();
-			const downloadResult = await this.download(
+			const status = await this.download(
 				client,
 				binSource,
 				writeStream,
@@ -500,7 +486,7 @@ export class CliManager {
 				onProgress,
 			);
 
-			switch (downloadResult.status) {
+			switch (status) {
 				case 200: {
 					await downloadProgress.writeProgress(progressLogPath, {
 						bytesDownloaded: 0,
@@ -531,14 +517,11 @@ export class CliManager {
 					// Replace existing binary (handles both renames + Windows lock)
 					await this.replaceExistingBinary(binPath, tempFile);
 
-					return {
-						binPath,
-						downloadedBytes: downloadResult.downloadedBytes,
-					};
+					return binPath;
 				}
 				case 304: {
 					this.output.info("Using existing binary since server returned a 304");
-					return { binPath };
+					return binPath;
 				}
 				case 404: {
 					vscode.window
@@ -575,7 +558,7 @@ export class CliManager {
 							}
 							const params = new URLSearchParams({
 								title: `Failed to download binary on \`${cliUtils.goos()}-${cliUtils.goarch()}\``,
-								body: `Received status code \`${downloadResult.status}\` when downloading the binary.`,
+								body: `Received status code \`${status}\` when downloading the binary.`,
 							});
 							const uri = vscode.Uri.parse(
 								`https://github.com/coder/vscode-coder/issues/new?${params.toString()}`,
@@ -603,7 +586,7 @@ export class CliManager {
 			bytesDownloaded: number,
 			totalBytes: number | null,
 		) => Promise<void>,
-	): Promise<DownloadResult> {
+	): Promise<number> {
 		const baseUrl = client.defaults.baseURL;
 
 		const controller = new AbortController();
@@ -621,8 +604,8 @@ export class CliManager {
 		});
 		this.output.info("Got status code", resp.status);
 
-		let written = 0;
 		if (resp.status === 200) {
+			let written = 0;
 			const rawContentLength = (resp.headers["content-length"] ??
 				resp.headers["x-original-content-length"]) as unknown;
 			const contentLength = Number.parseInt(
@@ -722,10 +705,7 @@ export class CliManager {
 			this.output.info(`Downloaded ${prettyBytes(written)}`);
 		}
 
-		return {
-			status: resp.status,
-			downloadedBytes: written,
-		};
+		return resp.status;
 	}
 
 	/**
@@ -815,8 +795,8 @@ export class CliManager {
 		this.output.info("Downloading signature from", source);
 		const signaturePath = path.join(cliPath + ".asc");
 		const writeStream = createWriteStream(signaturePath);
-		const downloadResult = await this.download(client, source, writeStream);
-		if (downloadResult.status === 200) {
+		const status = await this.download(client, source, writeStream);
+		if (status === 200) {
 			try {
 				await pgp.verifySignature(
 					publicKeys,
@@ -845,7 +825,7 @@ export class CliManager {
 				this.output.info("Binary will be ran anyway at user request");
 			}
 		}
-		return downloadResult.status;
+		return status;
 	}
 
 	/**

src/extension.ts

@@ -14,7 +14,10 @@ import { ServiceContainer } from "./core/container";
 import { DeploymentManager } from "./deployment/deploymentManager";
 import { CertificateError } from "./error/certificateError";
 import { getErrorDetail, toError } from "./error/errorUtils";
-import { ActivationTelemetry } from "./instrumentation/activation";
+import {
+	ActivationTelemetry,
+	type ActivationTracer,
+} from "./instrumentation/activation";
 import { OAuthSessionManager } from "./oauth/sessionManager";
 import { Remote } from "./remote/remote";
 import { getRemoteSshExtension } from "./remote/sshExtension";
@@ -37,15 +40,15 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 		serviceContainer.getTelemetryService(),
 	);
 
-	await activationTelemetry.trace(() =>
-		doActivate(ctx, serviceContainer, activationTelemetry),
+	await activationTelemetry.trace((tracer) =>
+		doActivate(ctx, serviceContainer, tracer),
 	);
 }
 
 async function doActivate(
 	ctx: vscode.ExtensionContext,
 	serviceContainer: ServiceContainer,
-	activationTelemetry: ActivationTelemetry,
+	tracer: ActivationTracer,
 ): Promise<void> {
 	// The Remote SSH extension's proposed APIs are used to override the SSH host
 	// name in VS Code itself. It's visually unappealing having a lengthy name!
@@ -92,9 +95,7 @@ async function doActivate(
 	const deploymentSessionAuth = deployment
 		? await secretsManager.getSessionAuth(deployment.safeHostname)
 		: undefined;
-	activationTelemetry.setAuthState(
-		deploymentSessionAuth ? "valid_token" : "none",
-	);
+	tracer.setAuthState(deploymentSessionAuth ? "valid_token" : "none");
 
 	// Shared handler for auth failures (used by interceptor + session manager)
 	const handleAuthFailure = (): Promise<void> => {
@@ -396,9 +397,7 @@ async function doActivate(
 					url: details.url,
 					token: details.token,
 				});
-				activationTelemetry.setAuthState(
-					deploymentSet ? "valid_token" : "expired",
-				);
+				tracer.setAuthState(deploymentSet ? "valid_token" : "expired");
 
 				// If a deep link stored a chat agent ID before the
 				// remote-authority reload, open it now that the
@@ -456,7 +455,7 @@ async function doActivate(
 		contextManager.set("coder.loaded", true);
 	} else if (deployment) {
 		output.info(`Initializing deployment: ${deployment.url}`);
-		activationTelemetry
+		tracer
 			.traceDeploymentInit(() =>
 				deploymentManager.setDeploymentIfValid(deployment),
 			)

src/instrumentation/activation.ts

@@ -1,38 +1,34 @@
 import type { TelemetryService } from "../telemetry/service";
-import type { Span } from "../telemetry/span";
 
 export type ActivationAuthState = "none" | "valid_token" | "expired";
 
-export class ActivationTelemetry {
-	private authState: ActivationAuthState = "none";
-	private span: Span | undefined;
+/** Helpers scoped to the activation trace's lifetime. */
+export interface ActivationTracer {
+	setAuthState(state: ActivationAuthState): void;
+	traceDeploymentInit(fn: () => Promise<boolean>): Promise<boolean>;
+}
 
+export class ActivationTelemetry {
 	public constructor(private readonly telemetry: TelemetryService) {}
 
-	public trace<T>(fn: () => Promise<T>): Promise<T> {
+	public trace<T>(fn: (tracer: ActivationTracer) => Promise<T>): Promise<T> {
 		return this.telemetry.trace("activation", async (span) => {
-			this.span = span;
-			span.setProperty("authState", this.authState);
-			try {
-				return await fn();
-			} finally {
-				this.span = undefined;
-			}
-		});
-	}
-
-	public setAuthState(authState: ActivationAuthState): void {
-		this.authState = authState;
-		this.span?.setProperty("authState", authState);
-	}
-
-	public traceDeploymentInit(fn: () => Promise<boolean>): Promise<boolean> {
-		const initialAuthState = this.authState;
-		return this.telemetry.trace("activation.deployment_init", async (span) => {
-			span.setProperty("authState", initialAuthState);
-			const success = await fn();
-			span.setProperty("authState", success ? "valid_token" : "expired");
-			return success;
+			span.setProperty("authState", "none");
+			return fn({
+				setAuthState: (state) => span.setProperty("authState", state),
+				traceDeploymentInit: (initFn) =>
+					this.telemetry.trace(
+						"activation.deployment_init",
+						async (childSpan) => {
+							const success = await initFn();
+							childSpan.setProperty(
+								"authState",
+								success ? "valid_token" : "expired",
+							);
+							return success;
+						},
+					),
+			});
 		});
 	}
 }

src/instrumentation/remoteSetup.ts

@@ -1,4 +1,4 @@
-import type { Span } from "../telemetry/span";
+import type { TelemetryService } from "../telemetry/service";
 
 export type RemoteSetupPhase =
 	| "auth_retrieval"
@@ -7,13 +7,20 @@ export type RemoteSetupPhase =
 	| "agent_ready"
 	| "ssh_config_write";
 
+/** Helpers scoped to the remote.setup trace's lifetime. */
+export interface RemoteSetupTracer {
+	phase<T>(name: RemoteSetupPhase, fn: () => T | PromiseLike<T>): Promise<T>;
+}
+
 export class RemoteSetupTelemetry {
-	public constructor(private readonly span: Span) {}
+	public constructor(private readonly telemetry: TelemetryService) {}
 
-	public phase<T>(
-		phaseName: RemoteSetupPhase,
-		fn: () => T | PromiseLike<T>,
-	): Promise<T> {
-		return this.span.phase(phaseName, () => Promise.resolve(fn()));
+	public trace<T>(fn: (tracer: RemoteSetupTracer) => Promise<T>): Promise<T> {
+		return this.telemetry.trace("remote.setup", (span) =>
+			fn({
+				phase: (name, phaseFn) =>
+					span.phase(name, () => Promise.resolve(phaseFn())),
+			}),
+		);
 	}
 }