Log AHP frames for the Web PubSub relay transport

src/vs/platform/agentHost/browser/webPubSubRelayTransport.ts

@@ -16,6 +16,7 @@ import { Emitter } from '../../../base/common/event.js';
 import { Disposable, DisposableStore } from '../../../base/common/lifecycle.js';
 import { IntervalTimer, disposableTimeout } from '../../../base/common/async.js';
 import { AgentHostClientConnectionKind } from '../common/agentHostTelemetry.js';
+import { AhpJsonlLogger, getAhpLogByteLength } from '../common/ahpJsonlLogger.js';
 import type { AhpServerNotification, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, ProtocolMessage } from '../common/state/sessionProtocol.js';
 import type { IClientTransport } from '../common/state/sessionTransport.js';
 import { Reassembler } from '../common/webPubSub/chunking.js';
@@ -82,6 +83,12 @@ export interface IWebPubSubRelayTransportOptions {
 	readonly webSocketFactory?: WebSocketFactory;
 	/** Invoked when an inbound frame can't be parsed or framed. */
 	readonly onProtocolError?: (err: unknown) => void;
+	/**
+	 * Records every AHP frame to a JSONL transcript when
+	 * `chat.agentHost.ahpJsonlLoggingEnabled` is on. Cloud sandbox hosts do not implement
+	 * `vscode/collectAgentHostDebugLogs`, so this is the only way to see their frames.
+	 */
+	readonly ahpLogger?: AhpJsonlLogger;
 }
 
 /**
@@ -118,6 +125,9 @@ export class WebPubSubRelayTransport extends Disposable implements IClientTransp
 
 	constructor(private readonly _options: IWebPubSubRelayTransportOptions) {
 		super();
+		if (this._options.ahpLogger) {
+			this._register(this._options.ahpLogger);
+		}
 	}
 
 	get isOpen(): boolean {
@@ -257,7 +267,9 @@ export class WebPubSubRelayTransport extends Disposable implements IClientTransp
 			return;
 		}
 		if (result.kind === 'payload') {
-			this._onMessage.fire(result.payload as ProtocolMessage);
+			const payload = result.payload as ProtocolMessage;
+			this._options.ahpLogger?.log(payload, 's2c', getAhpLogByteLength(JSON.stringify(payload)));
+			this._onMessage.fire(payload);
 		}
 	}
 
@@ -273,6 +285,9 @@ export class WebPubSubRelayTransport extends Disposable implements IClientTransp
 		if (this._closed || !this._ws) {
 			throw new Error('WebPubSubRelayTransport is closed');
 		}
+		// Logged before chunking, so the transcript carries whole AHP messages rather than the
+		// relay frames they were split into.
+		this._options.ahpLogger?.log(message, 'c2s', getAhpLogByteLength(JSON.stringify(message)));
 		const frames = buildPublish({
 			group: this._options.toHostGroup,
 			nextAckId: () => ++this._ackId,

src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostService.ts

@@ -11,7 +11,8 @@ import { IProtocolTransport } from '../../../../../platform/agentHost/common/sta
 import { RemoteAgentHostProtocolClient } from '../../../../../platform/agentHost/browser/remoteAgentHostProtocolClient.js';
 import { editorWindowAgentHostClientInfo } from '../../../../../platform/agentHost/common/agentHostClientInfo.js';
 import { WebPubSubRelayTransport } from '../../../../../platform/agentHost/browser/webPubSubRelayTransport.js';
-import { GITHUB_COPILOT_PROTECTED_RESOURCE } from '../../../../../platform/agentHost/common/agentService.js';
+import { AhpJsonlLogger } from '../../../../../platform/agentHost/common/ahpJsonlLogger.js';
+import { GITHUB_COPILOT_PROTECTED_RESOURCE, AgentHostAhpJsonlLoggingSettingId } from '../../../../../platform/agentHost/common/agentService.js';
 import {
 	buildWpsUrl,
 	cloudSandboxAddress,
@@ -25,6 +26,7 @@ import {
 import { IRemoteAgentHostService, RemoteAgentHostConnectionStatus, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId } from '../../../../../platform/agentHost/common/remoteAgentHostService.js';
 import { PROTOCOL_VERSION } from '../../../../../platform/agentHost/common/state/protocol/version/registry.js';
 import { IConfigurationService } from '../../../../../platform/configuration/common/configuration.js';
+import { IEnvironmentService } from '../../../../../platform/environment/common/environment.js';
 import { IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js';
 import { ILogService } from '../../../../../platform/log/common/log.js';
 import { CloudSandboxCredentialRefresher, MAX_WAKING_DELAY_MS, type ICloudSandboxCreds } from './cloudSandboxCredentialRefresh.js';
@@ -57,6 +59,7 @@ export class CloudSandboxAgentHostService extends Disposable implements ICloudSa
 		@ICloudSandboxApiService private readonly _apiService: ICloudSandboxApiService,
 		@IConfigurationService private readonly _configurationService: IConfigurationService,
 		@IInstantiationService private readonly _instantiationService: IInstantiationService,
+		@IEnvironmentService private readonly _environmentService: IEnvironmentService,
 		@ILogService private readonly _logService: ILogService,
 	) {
 		super();
@@ -116,11 +119,20 @@ export class CloudSandboxAgentHostService extends Disposable implements ICloudSa
 		// Three per-client relay lanes: publish to `to_host`; receive replies on `to_client` and
 		// unsolicited session state on `broadcast`. `groupValidation` drops inbound frames whose
 		// group name doesn't carry our own client id.
+		// Each soft reconnect gets a transport-owned logger keyed by connection id.
+		const ahpLoggingEnabled = !!this._configurationService.getValue<boolean>(AgentHostAhpJsonlLoggingSettingId);
 		const transportFactory = (): IProtocolTransport => new WebPubSubRelayTransport({
 			url: buildWpsUrl(creds.token),
 			toHostGroup: creds.token.groups.to_host,
 			joinGroups: [creds.token.groups.broadcast, creds.token.groups.to_client],
 			groupValidation: { expected: { cid: creds.token.client_id } },
+			ahpLogger: ahpLoggingEnabled
+				? this._instantiationService.createInstance(AhpJsonlLogger, {
+					logsHome: this._environmentService.logsHome,
+					connectionId: clientToken.client_id,
+					transport: 'webpubsub',
+				})
+				: undefined,
 		});
 
 		// Mission Control mints the client id and binds the relay lane to it, so the AHP identity