@codingame/monaco-vscode-api 37.6.0 → 38.0.0

package/vscode/src/vs/platform/agentHost/browser/agentHostProtocolClient.js

@@ -1,21 +1,28 @@
 
 import { __decorate, __param } from '../../../../../../external/tslib/tslib.es6.js';
-import { TimeoutTimer, DeferredPromise } from '../../../base/common/async.js';
+import { TimeoutTimer, IntervalTimer, DeferredPromise, disposableTimeout } from '../../../base/common/async.js';
 import { CancellationError } from '../../../base/common/errors.js';
 import { Emitter, Event } from '../../../base/common/event.js';
-import { Disposable, MutableDisposable, DisposableStore } from '../../../base/common/lifecycle.js';
+import { Disposable, MutableDisposable, DisposableStore, toDisposable } from '../../../base/common/lifecycle.js';
 import { Schemas } from '../../../base/common/network.js';
 import { hasKey } from '../../../base/common/types.js';
 import { URI } from '../../../base/common/uri.js';
 import { generateUuid } from '../../../base/common/uuid.js';
+import { newAgentHostSessionUri } from '../common/agentHostSessionIdentity.js';
+import { AgentHostSessionUrisCapabilityMetaKey } from '../common/meta/vscode/agentHostSessionUrisMeta.js';
+import { readRelayKeepAliveTimeout } from '../common/meta/relayConnectionMeta.js';
+import { vObj, vOptionalProp, vString, vArray, vEnum } from '../../../base/common/validation.js';
 import { localize } from '../../../nls.js';
 import { ILogService } from '../../log/common/log.service.js';
 import { toFileSystemProviderErrorCode, FileSystemProviderErrorCode } from '../../files/common/files.js';
 import { ConfigurationTarget, ConfigurationTargetToString } from '../../configuration/common/configuration.js';
 import { IConfigurationService } from '../../configuration/common/configuration.service.js';
-import { AgentSession } from '../common/agent.js';
 import { AGENT_HOST_DEBUG_LOGS_MAX_ENTRIES, AGENT_HOST_DEBUG_LOGS_CHUNK_BYTES } from '../common/agentService.js';
-import { RemoveSessionArtifactExtensionMethod, CreateAgentHostDetachedWorktreeExtensionMethod, SetAgentHostDetachedWorktreeArchivedExtensionMethod, ClaimAgentHostDetachedWorktreeExtensionMethod, DeleteAgentHostDetachedWorktreeExtensionMethod, ReconcileAgentHostDetachedWorktreesExtensionMethod, supportsAgentHostChatStateFile, GetAgentHostSessionStateFileExtensionMethod, CollectAgentHostDebugLogsExtensionMethod, ReadAgentHostDebugLogsChunkExtensionMethod, RequestAgentHostWorkspaceTrustExtensionMethod } from '../common/agentHostExtensionProtocol.js';
+import { supportsAgentHostDebugLogsCancellation } from '../common/meta/vscode/agentHostDebugLogsMeta.js';
+import { ReportAgentHostFirstResponseExtensionMethod, ReportChatUserInteractionExtensionMethod, RemoveSessionArtifactExtensionMethod, StopBackgroundWorkExtensionMethod, ImportSessionExtensionMethod, CreateAgentHostDetachedWorktreeExtensionMethod, SetAgentHostDetachedWorktreeArchivedExtensionMethod, ClaimAgentHostDetachedWorktreeExtensionMethod, DeleteAgentHostDetachedWorktreeExtensionMethod, ReconcileAgentHostDetachedWorktreesExtensionMethod, supportsAgentHostChatStateFile, GetAgentHostSessionStateFileExtensionMethod, CollectAgentHostDebugLogsExtensionMethod, CancelAgentHostDebugLogsExtensionMethod, ReadAgentHostDebugLogsChunkExtensionMethod, RequestAgentHostWorkspaceTrustExtensionMethod, RequestAgentHostMcpAuthenticationExtensionMethod } from '../common/agentHostExtensionProtocol.js';
+import { McpAuthRequiredReason } from '../common/state/protocol/channels-session/state.js';
+import { supportsAgentHostTiming, supportsChatUserInteractionTiming } from '../common/meta/vscode/agentHostTimingMeta.js';
+import { readCodexSessionModel } from '../common/meta/codexSessionModel.js';
 import { AMBIENT_AGENT_HOST_AUTHORITY } from '../common/agentHostConnectionsService.js';
 import { createRemoteWatchHandle } from '../common/agentHostFileSystemProvider.js';
 import { AgentSubscriptionManager } from '../common/state/agentSubscription.js';
@@ -24,10 +31,10 @@ import { LOCAL_AGENT_HOST_RESOURCE_IDENTITY, AgentHostResourcePermissionError }
 import { IAgentHostResourceService } from '../common/agentHostResourceService.service.js';
 import { ActionType } from '../common/state/protocol/common/actions.js';
 import '../common/state/protocol/common/notifications.js';
-import { ROOT_STATE_URI, isAhpRootChannel, StateComponents, isDefaultChatUri } from '../common/state/sessionState.js';
+import { ROOT_STATE_URI, StateComponents, isAhpRootChannel, isDefaultChatUri, isSessionStatusRead, isSessionChatArchived } from '../common/state/sessionState.js';
 import { normalizeLegacyActionEnvelope } from '../common/state/legacyProtocolCompatibility.js';
-import { SUPPORTED_PROTOCOL_VERSIONS } from '../common/state/protocol/version/registry.js';
-import { ProtocolError, isJsonRpcRequest, isJsonRpcResponse, isJsonRpcNotification } from '../common/state/sessionProtocol.js';
+import { getAgentHostSupportedProtocolVersions } from '../common/agentHostProtocolCompatibility.js';
+import { ProtocolError, JSON_RPC_INTERNAL_ERROR, isJsonRpcRequest, isJsonRpcResponse, isJsonRpcNotification } from '../common/state/sessionProtocol.js';
 import { isClientTransport, NonReconnectableTransportError } from '../common/state/sessionTransport.js';
 import { AhpErrorCodes, JsonRpcErrorCodes } from '../common/state/protocol/common/errors.js';
 import { ReconnectResultType, ContentEncoding } from '../common/state/protocol/common/commands.js';
@@ -41,26 +48,29 @@ import { TELEMETRY_SETTING_ID, TELEMETRY_OLD_SETTING_ID, TELEMETRY_CRASH_REPORTE
 import { ITelemetryService } from '../../telemetry/common/telemetry.service.js';
 import { getTelemetryLevel } from '../../telemetry/common/telemetryUtils.js';
 import { GLOBAL_AUTO_APPROVE_SETTING_ID, TERMINAL_AUTO_APPROVE_ENABLED_SETTING_ID, TERMINAL_AUTO_APPROVE_SETTING_ID, TERMINAL_IGNORE_DEFAULT_AUTO_APPROVE_RULES_SETTING_ID, DISABLE_REPO_INFO_TELEMETRY_SETTING_ID, AgentHostAutoApprovePolicyRestrictedConfigKey, AgentHostWorkspaceTrustConfigKey, AgentHostTelemetryLevelConfigKey, telemetryLevelToAgentHostConfigValue, AgentHostDisableRepoInfoTelemetryConfigKey, AgentHostTerminalAutoApproveEnabledConfigKey, AgentHostTerminalAutoApproveRulesConfigKey, getAgentHostTerminalAutoApproveRulesConfig } from '../common/agentHostSchema.js';
-import { getAgentHostConfigurationSyncEntries, getAgentHostConfigurationSyncTarget, resolveAgentHostConfigurationSyncValue, formatAgentHostConfigurationSyncValueForLog, resolveAgentHostConfigurationSyncPatch } from '../common/agentHostConfigurationSync.js';
+import { getAgentHostConfigurationSyncEntries, getAgentHostConfigurationSyncTarget, resolveAgentHostConfigurationSyncValue, addAgentHostConfigurationSyncValue, formatAgentHostConfigurationSyncValueForLog, resolveAgentHostConfigurationSyncPatch } from '../common/agentHostConfigurationSync.js';
 import { managedPermissionsConfigurationIds, resolveManagedSettingsPermissions } from '../common/agentHostManagedSettings.js';
 import { toAgentHostClientMeta, AgentHostClientConnectionKind } from '../common/agentHostTelemetry.js';
 import '../../../base/common/observableInternal/index.js';
 import { isFileResourceRead } from '../common/resourceReadLogging.js';
 import { ResourceSet } from '../../../base/common/map.js';
+import { StopWatch } from '../../../base/common/stopwatch.js';
 import { DEFAULT_RECONNECT_POLICY, hasExhaustedReconnectAttempts, computeReconnectDelay } from '../common/reconnectPolicy.js';
 import { IWorkspaceTrustEnablementService, IWorkspaceTrustManagementService, IWorkspaceTrustRequestService } from '../../workspace/common/workspaceTrust.service.js';
 import { isWorktreeUnderRepository } from '../common/worktreePaths.js';
+import { traceConnectionOperation, sanitizeConnectionDiagnosticText, getConnectionDiagnosticError } from '../common/connectionDiagnostics.js';
+import { DevContainerAgentHostProtocolClient } from '../common/devContainerAgentHostProtocolClient.js';
+import { supportsAgentHostDevContainers } from '../common/meta/vscode/agentHostDevContainersMeta.js';
 import { observableValue } from '../../../base/common/observableInternal/observables/observableValue.js';
 import { MessageAttachmentKind } from '../common/state/protocol/channels-chat/state.js';
 
 const AHP_CLIENT_CONNECTION_CLOSED = -32e3;
-const CLIENT_SUPPORTED_PROTOCOL_VERSIONS = SUPPORTED_PROTOCOL_VERSIONS.filter(version => version !== "0.8.0");
 const PING_INTERVAL_MS = 5_000;
 const LIVENESS_TIMEOUT_MS = 20_000;
 function connectionTimeoutError(address, silenceMs) {
     return ( new ProtocolError(
         AHP_CLIENT_CONNECTION_CLOSED,
-        `Connection appears dead: ${address}; no message received for ${silenceMs}ms.`
+        `Connection appears dead: ${address}; no data received for ${silenceMs}ms.`
     ));
 }
 function connectionClosedError(address) {
@@ -91,7 +101,52 @@ class InitialAuthenticationError extends Error {
         this.name = "InitialAuthenticationError";
     }
 }
+const mcpAuthenticationRequestValidator = vObj({
+    serverName: vString(),
+    auth: vObj({
+        reason: vEnum(
+            McpAuthRequiredReason.Required,
+            McpAuthRequiredReason.Expired,
+            McpAuthRequiredReason.InsufficientScope
+        ),
+        resource: vObj({
+            resource: vString(),
+            resource_name: vOptionalProp(vString()),
+            authorization_servers: vOptionalProp(vArray(vString())),
+            scopes_supported: vOptionalProp(vArray(vString()))
+        }),
+        oauthClient: vOptionalProp(vObj({
+            clientId: vString(),
+            clientSecret: vOptionalProp(vString())
+        })),
+        requiredScopes: vOptionalProp(vArray(vString())),
+        description: vOptionalProp(vString())
+    })
+});
 let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
+    get clientConnectionKind() {
+        return this._transport.clientConnectionKind;
+    }
+    get devContainerService() {
+        return this._state.kind !== AgentHostClientState.Closed && this._state.kind !== AgentHostClientState.Incompatible && supportsAgentHostDevContainers(this._initializeResult.get()) ? this._devContainerService : undefined;
+    }
+    _diagnostic(phase, detail) {
+        this._onDidConnectionDiagnostic.fire({
+            operationId: this._clientId,
+            attemptId: this._diagnosticAttemptId,
+            phase,
+            outcome: "info",
+            timestamp: Date.now(),
+            detail
+        });
+    }
+    _traceConnection(phase, operation) {
+        const attemptId = this._diagnosticAttemptId;
+        return traceConnectionOperation(event => this._onDidConnectionDiagnostic.fire({
+            ...event,
+            attemptId
+        }), phase, operation);
+    }
     get clientId() {
         return this._clientId;
     }
@@ -134,6 +189,7 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         this._serverSeq = 0;
         this._nextClientSeq = 1;
         this._initializeResult = observableValue("agentHostInitializeResult", undefined);
+        this._devContainerService = this._register(( new DevContainerAgentHostProtocolClient((method, params) => this._sendExtensionRequest(method, params))));
         this._onDidAction = this._register(( new Emitter()));
         this.onDidAction = this._onDidAction.event;
         this._onDidNotification = this._register(( new Emitter()));
@@ -148,8 +204,13 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         this.onDidFatalClose = this._onDidFatalClose.event;
         this._onDidChangeConnectionState = this._register(( new Emitter()));
         this.onDidChangeConnectionState = this._onDidChangeConnectionState.event;
+        this._onDidReconnect = this._register(( new Emitter()));
+        this.onDidReconnect = this._onDidReconnect.event;
         this._onDidScheduleReconnect = this._register(( new Emitter()));
         this.onDidScheduleReconnect = this._onDidScheduleReconnect.event;
+        this._onDidConnectionDiagnostic = this._register(( new Emitter()));
+        this.onDidConnectionDiagnostic = this._onDidConnectionDiagnostic.event;
+        this._diagnosticAttemptId = generateUuid();
         this._state = {
             kind: AgentHostClientState.Connecting,
             outbox: []
@@ -163,7 +224,12 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         this._authenticationRestorePending = false;
         this._lastReadTime = Date.now();
         this._pingTimer = this._register(new TimeoutTimer());
+        this._relayKeepAliveTimer = this._register(( new IntervalTimer()));
         this._closeTimer = this._register(new TimeoutTimer());
+        this._reconnectDeadlineTimer = this._register(new TimeoutTimer());
+        this._isSuspended = false;
+        this._firstSessionRequestPending = false;
+        this._didDispose = false;
         this._grantedImplicitReadUris = ( new ResourceSet());
         this._implicitReadGrants = this._register(( new DisposableStore()));
         this._resourceIdentity = identity;
@@ -172,9 +238,16 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         this._connectionAuthority = identity === LOCAL_AGENT_HOST_RESOURCE_IDENTITY ? AMBIENT_AGENT_HOST_AUTHORITY : agentHostAuthority(identity);
         this.resourceUris = identity === LOCAL_AGENT_HOST_RESOURCE_IDENTITY ? identityAgentHostResourceUriMapper : createAgentHostResourceUriMapper(this._connectionAuthority);
         this._loadEstimator = options?.loadEstimator ?? LoadEstimator.getInstance();
+        if (options?.onDidSuspend && options.onDidResume) {
+            this._register(options.onDidSuspend(() => this._suspendReconnectDeadline()));
+            this._register(options.onDidResume(() => this._resumeReconnectDeadline()));
+        }
         this._clientInfo = options?.clientInfo;
         this._reconnectPolicy = options?.reconnectPolicy ?? DEFAULT_RECONNECT_POLICY;
+        this._prepareReconnect = options?.prepareReconnect;
+        this._prepareAuthentication = options?.prepareAuthentication;
         this._resolveInitialAuthentication = options?.resolveInitialAuthentication;
+        this._onDispose = options?.onDispose;
         if (typeof transportOrFactory === "function") {
             this._transportFactory = transportOrFactory;
             this._installTransport(transportOrFactory());
@@ -193,7 +266,7 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
             this._subscriptionManager.receiveEnvelope(envelope);
         }));
         this._register(this._configurationService.onDidChangeConfiguration(e => {
-            if (this._state.kind !== AgentHostClientState.Connected) {
+            if (this._state.kind !== AgentHostClientState.Connected || this._isWebPubSubRelay()) {
                 return;
             }
             const patch = {};
@@ -205,10 +278,11 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
                     }
                     const value = resolveAgentHostConfigurationSyncValue(this._configurationService, entry);
                     if (value !== undefined) {
-                        patch[entry.sync.key] = value;
-                        mirrored.push(
-                            `${entry.sync.key}=${formatAgentHostConfigurationSyncValueForLog(entry.settingId, value)} (${entry.settingId})`
-                        );
+                        for (const key of addAgentHostConfigurationSyncValue(patch, entry, value)) {
+                            mirrored.push(
+                                `${key}=${formatAgentHostConfigurationSyncValueForLog(entry.settingId, patch[key])} (${entry.settingId})`
+                            );
+                        }
                     }
                 }
             }
@@ -253,7 +327,22 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         const listeners = ( new DisposableStore());
         listeners.add(transport);
         listeners.add(transport.onMessage(msg => this._handleMessage(msg)));
+        if (transport.onDidReceiveData) {
+            listeners.add(transport.onDidReceiveData(() => this._recordInboundTraffic()));
+        }
         listeners.add(transport.onClose(() => this._handleTransportClose()));
+        if (transport.onDidCloseDetails) {
+            listeners.add(transport.onDidCloseDetails(close => {
+                this._diagnostic(
+                    "transport.closeDetails",
+                    `code=${close.code ?? "unavailable"}; wasClean=${close.wasClean ?? "unavailable"}; reason=${sanitizeConnectionDiagnosticText(close.reason ?? "(unavailable)")}`
+                );
+            }));
+        }
+        if (this._relayRootRefresh && this._relayRootRefresh.transport !== transport) {
+            this._subscriptionManager.cancelSnapshotRefresh(( URI.parse(ROOT_STATE_URI)));
+            this._relayRootRefresh = undefined;
+        }
         this._transport = transport;
         this._transportListeners.value = listeners;
     }
@@ -262,6 +351,7 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
             return;
         }
         if (this._state.kind === AgentHostClientState.Reconnecting) {
+            this._reconnectDeadlineTimer.cancel();
             const reconnect = this._state.reconnect;
             if (reconnect.timeoutHandle !== undefined) {
                 clearTimeout(reconnect.timeoutHandle);
@@ -270,6 +360,15 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
             reconnect.nextAttemptAt = undefined;
         }
         this._state = next;
+        if (next.kind === AgentHostClientState.Connected) {
+            this._firstSessionRequestPending = true;
+        }
+        if (next.kind === AgentHostClientState.Reconnecting) {
+            this._armReconnectDeadline(next.reconnect);
+        }
+        if (next.kind === AgentHostClientState.Reconnecting || next.kind === AgentHostClientState.Closed) {
+            this._devContainerService.connectionClosed();
+        }
         this._onDidChangeConnectionState.fire(next.kind);
     }
     _newReconnectGate() {
@@ -282,44 +381,115 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
             gate: this._newReconnectGate(),
             outbox: [],
             attempt: 0,
+            transportConnected: false,
             timeoutHandle: undefined,
-            nextAttemptAt: undefined
+            nextAttemptAt: undefined,
+            deadline: this._isSuspended || this._reconnectPolicy.maxElapsedTimeMs === undefined ? undefined : Date.now() + this._reconnectPolicy.maxElapsedTimeMs,
+            remainingTimeMs: this._isSuspended ? this._reconnectPolicy.maxElapsedTimeMs : undefined
         };
     }
+    _armReconnectDeadline(reconnect) {
+        if (reconnect.deadline !== undefined) {
+            this._reconnectDeadlineTimer.setIfNotSet(() => {
+                if (this._state.kind === AgentHostClientState.Reconnecting && this._state.reconnect === reconnect) {
+                    this._handleReconnectTimeout();
+                }
+            }, Math.max(0, reconnect.deadline - Date.now()));
+        }
+    }
+    _suspendReconnectDeadline() {
+        this._isSuspended = true;
+        if (this._state.kind === AgentHostClientState.Reconnecting) {
+            const reconnect = this._state.reconnect;
+            if (reconnect.deadline !== undefined) {
+                reconnect.remainingTimeMs = Math.max(0, reconnect.deadline - Date.now());
+                reconnect.deadline = undefined;
+                this._reconnectDeadlineTimer.cancel();
+            }
+        }
+    }
+    _resumeReconnectDeadline() {
+        this._isSuspended = false;
+        if (this._state.kind === AgentHostClientState.Reconnecting) {
+            const reconnect = this._state.reconnect;
+            if (reconnect.remainingTimeMs !== undefined) {
+                reconnect.deadline = Date.now() + reconnect.remainingTimeMs;
+                reconnect.remainingTimeMs = undefined;
+                this._armReconnectDeadline(reconnect);
+            }
+        }
+    }
+    _handleReconnectTimeout() {
+        const message = ( localize(
+            1862,
+            "Automatic reconnect timed out after {0}ms.",
+            this._reconnectPolicy.maxElapsedTimeMs
+        ));
+        this._diagnostic("reconnect.timedOut", message);
+        this._logService.warn(`[RemoteAgentHostProtocol] ${message}`);
+        this._handleFatalClose(( new ProtocolError(AHP_CLIENT_CONNECTION_CLOSED, message)));
+    }
+    _checkReconnectState(reconnect) {
+        if (this._state.kind !== AgentHostClientState.Reconnecting || this._state.reconnect !== reconnect) {
+            return false;
+        }
+        if (reconnect.deadline !== undefined && Date.now() >= reconnect.deadline) {
+            this._handleReconnectTimeout();
+            return false;
+        }
+        return true;
+    }
     dispose() {
+        if (!this._didDispose) {
+            this._didDispose = true;
+            try {
+                this._onDispose?.();
+            } catch (error) {
+                this._logService.warn("[AgentHostProtocolClient] Dispose callback failed", error);
+            }
+        }
         this._handleClose(connectionDisposedError(this._address));
         super.dispose();
     }
     async connect() {
+        const watch = StopWatch.create(false);
+        let stage = "transport";
         try {
             if (isClientTransport(this._transport)) {
-                await this._raceClose(this._transport.connect());
+                const transport = this._transport;
+                await this._traceConnection("transport.connect", () => this._raceClose(transport.connect()));
             }
             if (this._state.kind !== AgentHostClientState.Connecting) {
                 throw transportLostError(this._address);
             }
-            const result = await this._dispatchRequest("initialize", {
+            this._lastReadTime = Date.now();
+            this._resetLivenessTimers();
+            stage = "initialize";
+            const result = await this._traceConnection("protocol.initialize", () => this._dispatchRequest("initialize", {
                 channel: ROOT_STATE_URI,
-                protocolVersions: [...CLIENT_SUPPORTED_PROTOCOL_VERSIONS],
+                protocolVersions: getAgentHostSupportedProtocolVersions(),
                 clientId: this._clientId,
                 clientInfo: this._clientInfo,
                 _meta: this._clientMeta(),
                 initialSubscriptions: [ROOT_STATE_URI]
             }, {
                 bypassInitializeQueue: true
-            });
+            }));
+            this._logService.info(
+                `[RemoteAgentHostProtocol] Initialized: address=${this._address} clientId=${this._clientId} durationMs=${watch.elapsed()}`
+            );
             this._applyInitializeResult(result);
+            this._hydrateInitializedRootState(result);
             if (this._resolveInitialAuthentication || this._authentication.size > 0) {
-                await this._restoreAuthenticationAfterFreshInitialize(AgentHostClientState.Connecting);
+                stage = "authentication";
+                await this._traceConnection(
+                    "protocol.authentication",
+                    () => this._restoreAuthentication(AgentHostClientState.Connecting)
+                );
                 if (this._state.kind !== AgentHostClientState.Connecting) {
                     throw transportLostError(this._address);
                 }
             }
-            for (const snapshot of result.snapshots ?? []) {
-                if (isAhpRootChannel(snapshot.resource)) {
-                    this._subscriptionManager.handleRootSnapshot(snapshot.state, snapshot.fromSeq);
-                }
-            }
             if (isClientTransport(this._transport) && this._state.kind === AgentHostClientState.Connecting) {
                 for (const message of this._state.outbox) {
                     this._transport.send(message);
@@ -329,8 +499,12 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
             this._transitionTo({
                 kind: AgentHostClientState.Connected
             });
+            this._diagnostic("connection.ready", `clientId=${this._clientId}`);
             this._resetLivenessTimers();
         } catch (error) {
+            this._logService.warn(
+                `[RemoteAgentHostProtocol] Connection failed: address=${this._address} clientId=${this._clientId} stage=${stage} durationMs=${watch.elapsed()} pendingRequests=${this._pendingRequests.size}`
+            );
             const protocolError = error instanceof ProtocolError ? error : ( new ProtocolError(
                 AHP_CLIENT_CONNECTION_CLOSED,
                 error instanceof Error ? error.message : String(error)
@@ -354,7 +528,7 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
             if (this._state.kind === AgentHostClientState.Reconnecting) {
                 throw error;
             }
-            if (protocolError.code === AHP_CLIENT_CONNECTION_CLOSED && this._beginReconnectFromConnecting(protocolError)) {
+            if ((protocolError.code === AHP_CLIENT_CONNECTION_CLOSED || (stage === "authentication" && protocolError.code === JSON_RPC_INTERNAL_ERROR)) && this._beginReconnectFromConnecting(protocolError)) {
                 throw error;
             }
             this._handleClose(protocolError);
@@ -365,6 +539,13 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         this._handleTransportClose();
     }
     _handleTransportClose() {
+        this._relayKeepAliveTimer.cancel();
+        if (this._state.kind !== AgentHostClientState.Closed) {
+            this._diagnostic(
+                "transport.closed",
+                `state=${this._state.kind}; sinceLastDataMs=${Date.now() - this._lastReadTime}`
+            );
+        }
         switch (this._state.kind) {
         case AgentHostClientState.Closed:
             return;
@@ -406,6 +587,7 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
             this._logService.info(
                 `[RemoteAgentHostProtocol] Transport lost for ${this._address} mid-reconnect; aborting the current attempt.`
             );
+            this._state.reconnect.transportConnected = false;
             this._cancelLivenessTimers();
             this._rejectPendingRequests(transportLostError(this._address));
             return;
@@ -417,13 +599,13 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         }
         if (!this._reconnectPolicy.autoRestore) {
             this._logService.info(
-                `[RemoteAgentHostProtocol] Transport lost while connecting to ${this._address}; automatic reconnect is disabled.`
+                `[RemoteAgentHostProtocol] Connection setup failed for ${this._address}; automatic reconnect is disabled.`
             );
             this._handleFatalClose(error);
             return true;
         }
         this._logService.info(
-            `[RemoteAgentHostProtocol] Transport lost while connecting to ${this._address}; scheduling a fresh initialize.`
+            `[RemoteAgentHostProtocol] Connection setup failed for ${this._address}; scheduling a fresh initialize.`
         );
         const outbox = this._state.outbox;
         this._rejectPendingRequests(error);
@@ -481,6 +663,7 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
             return;
         }
         if (!userInitiated && hasExhaustedReconnectAttempts(this._reconnectPolicy, reconnect.attempt)) {
+            this._diagnostic("reconnect.exhausted", `attempts=${reconnect.attempt}`);
             const error = ( new ProtocolError(
                 AHP_CLIENT_CONNECTION_CLOSED,
                 `Automatic reconnect gave up after ${reconnect.attempt} attempts.`
@@ -493,6 +676,7 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         }
         const attempt = reconnect.attempt + 1;
         const delay = computeReconnectDelay(this._reconnectPolicy, attempt);
+        this._diagnostic("reconnect.scheduled", `attempt=${attempt}; delayMs=${delay}`);
         this._logService.info(
             `[RemoteAgentHostProtocol] Reconnecting to ${this._address} in ${delay}ms (attempt ${attempt}).`
         );
@@ -511,17 +695,36 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
             return;
         }
         const reconnect = this._state.reconnect;
+        if (!this._checkReconnectState(reconnect)) {
+            return;
+        }
         reconnect.attempt++;
+        const attempt = reconnect.attempt;
+        this._diagnosticAttemptId = generateUuid();
+        this._diagnostic(
+            "reconnect.started",
+            `attempt=${reconnect.attempt}; clientId=${this._clientId}`
+        );
         let transport;
         try {
+            if (this._prepareReconnect) {
+                await this._raceClose(this._prepareReconnect());
+                if (!this._checkReconnectState(reconnect) || reconnect.attempt !== attempt) {
+                    return;
+                }
+            }
             transport = this._transportFactory();
             this._installTransport(transport);
             if (isClientTransport(transport)) {
-                await transport.connect();
+                const clientTransport = transport;
+                await this._traceConnection("transport.reconnect", () => this._raceClose(clientTransport.connect()));
             }
-            if (this._state.kind !== AgentHostClientState.Reconnecting) {
+            if (!this._checkReconnectState(reconnect)) {
                 return;
             }
+            reconnect.transportConnected = true;
+            this._lastReadTime = Date.now();
+            this._resetLivenessTimers();
             const subscriptions = this._subscriptionManager.currentSubscriptionUris().map(u => ( u.toString()));
             if (!subscriptions.includes(ROOT_STATE_URI)) {
                 subscriptions.unshift(ROOT_STATE_URI);
@@ -530,24 +733,38 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
             const {
                 result,
                 freshInitialize
-            } = await this._reconnectOrInitialize(lastSeenServerSeq, subscriptions);
-            if (this._state.kind !== AgentHostClientState.Reconnecting) {
+            } = await this._traceConnection(
+                "protocol.reconnect",
+                () => this._reconnectOrInitialize(lastSeenServerSeq, subscriptions)
+            );
+            this._diagnostic(
+                "protocol.reconnect.result",
+                `mode=${freshInitialize ? "freshInitialize" : result.type}`
+            );
+            if (!this._checkReconnectState(reconnect)) {
                 return;
             }
-            this._applyReconnectResult(result, freshInitialize);
+            const restoreAuthentication = freshInitialize || !!this._resolveInitialAuthentication || this._authentication.size > 0 || this._authenticationRestorePending;
+            this._applyReconnectResult(result, freshInitialize, restoreAuthentication);
             this._updateManagedSettingsPermissions(true);
-            if ((freshInitialize && result.type === ReconnectResultType.Snapshot) || this._authenticationRestorePending) {
-                if (freshInitialize && result.type === ReconnectResultType.Snapshot) {
+            if (restoreAuthentication) {
+                if (result.type === ReconnectResultType.Snapshot) {
                     this._markSubscriptionsAwaitingRestore(result.snapshots);
                 }
-                await this._restoreAuthenticationAfterFreshInitialize(AgentHostClientState.Reconnecting);
-                if (this._state.kind !== AgentHostClientState.Reconnecting) {
+                await this._traceConnection(
+                    "protocol.authentication",
+                    () => this._restoreAuthentication(AgentHostClientState.Reconnecting)
+                );
+                if (!this._checkReconnectState(reconnect)) {
                     return;
                 }
             }
             if (this._subscriptionsAwaitingRestore.size > 0) {
-                await this._restoreSubscriptionsAwaitingRestore();
-                if (this._state.kind !== AgentHostClientState.Reconnecting) {
+                await this._traceConnection(
+                    "protocol.subscriptions",
+                    () => this._restoreSubscriptionsAwaitingRestore()
+                );
+                if (!this._checkReconnectState(reconnect)) {
                     return;
                 }
             }
@@ -561,14 +778,31 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
             this._transitionTo({
                 kind: AgentHostClientState.Connected
             });
+            this._startRelayKeepAlive(result._meta ?? this._initializeResult.get()?._meta);
             gate.complete();
+            this._onDidReconnect.fire();
             this._logService.info(`[RemoteAgentHostProtocol] Reconnected to ${this._address}.`);
+            this._diagnostic("reconnect.succeeded", `attempt=${reconnect.attempt}`);
         } catch (err) {
+            if (!this._checkReconnectState(reconnect) || reconnect.attempt !== attempt) {
+                transport?.dispose();
+                return;
+            }
+            this._onDidConnectionDiagnostic.fire({
+                operationId: this._clientId,
+                attemptId: this._diagnosticAttemptId,
+                phase: "reconnect",
+                outcome: "failed",
+                timestamp: Date.now(),
+                error: getConnectionDiagnosticError(err)
+            });
             this._logService.warn(
                 `[RemoteAgentHostProtocol] Reconnect attempt failed for ${this._address}: ${err instanceof Error ? err.message : String(err)}`
             );
+            reconnect.transportConnected = false;
+            this._cancelLivenessTimers();
             transport?.dispose();
-            if (this._state.kind !== AgentHostClientState.Reconnecting) {
+            if (!this._checkReconnectState(reconnect)) {
                 return;
             }
             if (err instanceof NonReconnectableTransportError) {
@@ -596,6 +830,7 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
     async _reconnectOrInitialize(lastSeenServerSeq, subscriptions) {
         try {
             const result = await this._dispatchRequest("reconnect", {
+                channel: ROOT_STATE_URI,
                 clientId: this._clientId,
                 lastSeenServerSeq,
                 subscriptions,
@@ -617,7 +852,7 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         );
         const initializeResult = await this._dispatchRequest("initialize", {
             channel: ROOT_STATE_URI,
-            protocolVersions: [...CLIENT_SUPPORTED_PROTOCOL_VERSIONS],
+            protocolVersions: getAgentHostSupportedProtocolVersions(),
             clientId: this._clientId,
             clientInfo: this._clientInfo,
             _meta: this._clientMeta(),
@@ -634,8 +869,8 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
             freshInitialize: true
         };
     }
-    _markSubscriptionsAwaitingRestore(initialSnapshots) {
-        const restored = ( new Set(initialSnapshots.map(snapshot => snapshot.resource)));
+    _markSubscriptionsAwaitingRestore(snapshots) {
+        const restored = ( new Set(snapshots.map(snapshot => snapshot.resource)));
         this._subscriptionsAwaitingRestore.clear();
         for (const subscription of this._subscriptionManager.getActiveSubscriptions()) {
             const resource = ( subscription.resource.toString());
@@ -680,7 +915,7 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
                         throw error;
                     }
                     this._logService.warn(
-                        `[AgentHostProtocolClient] Failed to restore subscription ${resource} after host restart: ${error instanceof Error ? error.message : String(error)}`
+                        `[AgentHostProtocolClient] Failed to restore subscription ${resource} after reconnect: ${error instanceof Error ? error.message : String(error)}`
                     );
                     this._subscriptionManager.cancelSnapshotRefresh(subscription.resource);
                     this._subscriptionManager.markSubscriptionsMissing([subscription.resource]);
@@ -698,12 +933,23 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
             ))
         ]);
     }
-    async _restoreAuthenticationAfterFreshInitialize(expectedState) {
+    async _restoreAuthentication(expectedState) {
+        const state = this._state;
+        const transport = this._transport;
         this._authenticationRestorePending = true;
+        if (this._prepareAuthentication) {
+            await this._raceClose(this._prepareAuthentication());
+            if (this._state !== state || this._state.kind !== expectedState) {
+                return;
+            }
+        }
         let initialAuthenticationKey;
         if (this._resolveInitialAuthentication) {
             try {
-                const initialAuthentication = await this._resolveInitialAuthentication();
+                const initialAuthentication = await this._raceClose(this._resolveInitialAuthentication());
+                if (this._state !== state || this._state.kind !== expectedState) {
+                    return;
+                }
                 if (initialAuthentication) {
                     const normalizedParams = this._normalizeAuthenticationParams(initialAuthentication);
                     initialAuthenticationKey = this._authenticationKey(normalizedParams);
@@ -714,13 +960,21 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
                     });
                 }
             } catch (error) {
+                if (this._state !== state || isConnectionClosedError(error)) {
+                    throw error;
+                }
                 throw ( new InitialAuthenticationError(error));
             }
-            if (this._state.kind !== expectedState) {
+        }
+        const authenticationKeys = [...( this._authentication.keys())];
+        const restoreAuthentication = async key => {
+            if (this._state !== state || this._transport !== transport) {
+                throw transportLostError(this._address);
+            }
+            const authentication = this._authentication.get(key);
+            if (!authentication) {
                 return;
             }
-        }
-        await Promise.all([...this._authentication.entries()].map(async ([key, authentication]) => {
             const now = Date.now();
             if (isExpired(authentication.expiresAt, now)) {
                 this._authentication.delete(key);
@@ -743,31 +997,48 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
                     bypassReconnectGate: true
                 });
             } catch (error) {
-                if (isConnectionClosedError(error)) {
+                if (isConnectionClosedError(error) || (key === initialAuthenticationKey && error instanceof ProtocolError && error.code === JSON_RPC_INTERNAL_ERROR)) {
                     throw error;
                 }
                 if (key === initialAuthenticationKey) {
                     throw ( new InitialAuthenticationError(error));
                 }
                 this._logService.warn(
-                    `[AgentHostProtocolClient] Failed to restore authentication for ${params.resource} after host restart: ${error instanceof Error ? error.message : String(error)}`
+                    `[AgentHostProtocolClient] Failed to restore authentication for ${params.resource}: ${error instanceof Error ? error.message : String(error)}`
                 );
             }
-        }));
+        };
+        if (this._isWebPubSubRelay()) {
+            if (initialAuthenticationKey !== undefined) {
+                await restoreAuthentication(initialAuthenticationKey);
+            }
+            for (const key of authenticationKeys) {
+                if (key !== initialAuthenticationKey) {
+                    await restoreAuthentication(key);
+                }
+            }
+        } else {
+            await Promise.all(authenticationKeys.map(restoreAuthentication));
+        }
+        await this._refreshRelayRootSnapshot();
         this._authenticationRestorePending = false;
     }
     _clientMeta() {
         const telemetryLevel = this._effectiveTelemetryLevel();
         const sendIdentity = telemetryLevel >= TelemetryLevel.USAGE;
-        return toAgentHostClientMeta(
-            this._transport.clientConnectionKind,
-            telemetryLevel,
-            sendIdentity ? this._telemetryService.machineId : undefined,
-            sendIdentity ? this._telemetryService.devDeviceId : undefined
-        );
+        return {
+            ...toAgentHostClientMeta(
+                this._transport.clientConnectionKind,
+                telemetryLevel,
+                sendIdentity ? this._telemetryService.machineId : undefined,
+                sendIdentity ? this._telemetryService.devDeviceId : undefined
+            ),
+            [AgentHostSessionUrisCapabilityMetaKey]: true
+        };
     }
     _applyInitializeResult(result, forwardClientConfig = true) {
         this._initializeResult.set(result, undefined);
+        this._startRelayKeepAlive(result._meta);
         this._serverSeq = result.serverSeq;
         if (result.defaultDirectory) {
             const directory = result.defaultDirectory;
@@ -777,7 +1048,17 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
             this._forwardClientConfig();
         }
     }
+    _hydrateInitializedRootState(result) {
+        for (const snapshot of result.snapshots ?? []) {
+            if (isAhpRootChannel(snapshot.resource)) {
+                this._subscriptionManager.handleRootSnapshot(snapshot.state, snapshot.fromSeq);
+            }
+        }
+    }
     _forwardClientConfig(includeManagedSettings = true) {
+        if (this._isWebPubSubRelay()) {
+            return;
+        }
         this._dispatchRootConfig(resolveAgentHostConfigurationSyncPatch(
             this._configurationService,
             getAgentHostConfigurationSyncTarget(this._resourceIdentity)
@@ -813,7 +1094,7 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
             }
         });
     }
-    _applyReconnectResult(result, preservePending = false) {
+    _applyReconnectResult(result, preservePending = false, restoreMissingSubscriptions = false) {
         if (result.type === ReconnectResultType.Replay) {
             let maxSeq = this._serverSeq;
             for (const envelope of result.actions) {
@@ -830,9 +1111,15 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
                 this._logService.info(
                     `[RemoteAgentHostProtocol] Server cannot resume ${result.missing.length} subscription(s) after reconnect.`
                 );
-                this._subscriptionManager.markSubscriptionsMissing(result.missing.map(u => ( URI.parse(u))));
-                for (const resource of result.missing) {
-                    this._subscriptionsAwaitingRestore.delete(resource);
+                if (restoreMissingSubscriptions) {
+                    for (const resource of result.missing) {
+                        this._subscriptionsAwaitingRestore.add(resource);
+                    }
+                } else {
+                    this._subscriptionManager.markSubscriptionsMissing(result.missing.map(u => ( URI.parse(u))));
+                    for (const resource of result.missing) {
+                        this._subscriptionsAwaitingRestore.delete(resource);
+                    }
                 }
             }
         } else {
@@ -906,6 +1193,56 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         const seq = this._subscriptionManager.dispatchOptimistic(channel, action);
         this.dispatchAction(channel, action, this._clientId, seq);
     }
+    async dispatchConfirmed(channel, subscription, action, token) {
+        if (token.isCancellationRequested) {
+            throw ( new CancellationError());
+        }
+        if (this._state.kind === AgentHostClientState.Closed || this._state.kind === AgentHostClientState.Incompatible) {
+            throw this._state.error;
+        }
+        if (this._state.kind === AgentHostClientState.Reconnecting) {
+            throw transportLostError(this._address);
+        }
+        const resource = ( URI.parse(channel));
+        const held = isAhpRootChannel(channel) ? this.rootState : this._subscriptionManager.getSubscriptionUnmanaged(resource);
+        if (held !== subscription) {
+            throw ( new Error(`No matching held subscription for ${channel}`));
+        }
+        const store = ( new DisposableStore());
+        const confirmed = ( new DeferredPromise());
+        let clientSeq;
+        store.add(token.onCancellationRequested(() => confirmed.error(( new CancellationError()))));
+        store.add(Event.once(Event.filter(
+            this.onDidAction,
+            envelope => envelope.channel === channel && clientSeq !== undefined && envelope.origin?.clientId === this._clientId && envelope.origin.clientSeq === clientSeq
+        ))(envelope => confirmed.complete(envelope)));
+        store.add(Event.once(Event.filter(
+            this.onDidChangeConnectionState,
+            state => state === AgentHostClientState.Closed || state === AgentHostClientState.Incompatible || state === AgentHostClientState.Reconnecting
+        ))(() => confirmed.error(transportLostError(this._address))));
+        if (subscription.onDidError) {
+            store.add(Event.once(subscription.onDidError)(error => confirmed.error(error)));
+        }
+        try {
+            if (subscription.value === undefined) {
+                const initialized = ( new DeferredPromise());
+                store.add(Event.once(
+                    Event.filter(subscription.onDidChange, () => subscription.value !== undefined)
+                )(() => initialized.complete()));
+                await Promise.race([initialized.p, confirmed.p]);
+            }
+            if (subscription.value instanceof Error) {
+                throw subscription.value;
+            }
+            if (!confirmed.isSettled) {
+                clientSeq = this._subscriptionManager.dispatchOptimistic(channel, action);
+                this.dispatchAction(channel, action, this._clientId, clientSeq);
+            }
+            return await confirmed.p;
+        } finally {
+            store.dispose();
+        }
+    }
     async subscribe(resource) {
         this._logService.trace(`[RemoteAgentHostProtocol] subscribe start: ${( resource.toString())}`);
         const result = await this._sendRequest("subscribe", {
@@ -940,7 +1277,13 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         if (!provider) {
             throw ( new Error("Cannot create remote agent host session without a provider."));
         }
-        const session = config?.session ?? AgentSession.uri(provider, generateUuid());
+        const root = this.rootState.value;
+        const session = config?.session ?? newAgentHostSessionUri(
+            provider,
+            generateUuid(),
+            this.initializeResult.get(),
+            root instanceof Error ? undefined : root
+        );
         if (config?.activeClient?.customizations) {
             this._grantImplicitReadsForCustomizations(config.activeClient.customizations);
         }
@@ -956,12 +1299,37 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         this._subscriptionManager.trackSessionCreate(session, promise);
         return promise;
     }
+    async reportFirstResponse(diagnostic) {
+        if (supportsAgentHostTiming(this.initializeResult.get())) {
+            await this._sendExtensionRequest(ReportAgentHostFirstResponseExtensionMethod, diagnostic);
+        }
+    }
+    async reportUserInteraction(timing) {
+        if (supportsChatUserInteractionTiming(this.initializeResult.get())) {
+            await this._sendExtensionRequest(ReportChatUserInteractionExtensionMethod, timing);
+        }
+    }
     async removeSessionArtifact(session, artifactId) {
         await this._sendExtensionRequest(RemoveSessionArtifactExtensionMethod, {
             session: ( session.toString()),
             artifactId
         });
     }
+    async stopBackgroundWork(chat, id) {
+        const result = await this._sendExtensionRequest(StopBackgroundWorkExtensionMethod, {
+            chat: ( chat.toString()),
+            id
+        });
+        return result.stopped;
+    }
+    refreshSubscription(resource) {
+        return this._subscriptionManager.refreshSubscription(resource);
+    }
+    async importSession(session) {
+        await this._sendExtensionRequest(ImportSessionExtensionMethod, {
+            session: ( session.toString())
+        });
+    }
     async createDetachedWorktree(session, prompt) {
         const result = await this._sendExtensionRequest(CreateAgentHostDetachedWorktreeExtensionMethod, {
             session: ( session.toString()),
@@ -1045,6 +1413,15 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
             return this._initializeResult.get()?.completionTriggerCharacters ?? [];
         }
     }
+    registerMcpAuthenticationHandler(handler) {
+        if (this._mcpAuthenticationHandler) {
+            throw ( new Error("MCP authentication handler is already registered"));
+        }
+        this._mcpAuthenticationHandler = handler;
+        return toDisposable(() => {
+            this._mcpAuthenticationHandler = undefined;
+        });
+    }
     async authenticate(params) {
         const normalizedParams = this._normalizeAuthenticationParams(params);
         const expiresAt = getExpirationTime(params.expiresIn);
@@ -1062,10 +1439,63 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         } else {
             this._authentication.delete(key);
         }
+        await this._refreshRelayRootSnapshot();
         return {
             authenticated: true
         };
     }
+    _refreshRelayRootSnapshot() {
+        if (!this._isWebPubSubRelay()) {
+            return Promise.resolve();
+        }
+        if (this._relayRootRefresh?.transport === this._transport) {
+            return this._relayRootRefresh.promise;
+        }
+        const transport = this._transport;
+        const refresh = {
+            transport,
+            promise: this._readRelayRootSnapshot(transport)
+        };
+        this._relayRootRefresh = refresh;
+        const release = () => {
+            if (this._relayRootRefresh === refresh) {
+                this._relayRootRefresh = undefined;
+            }
+        };
+        void refresh.promise.then(release, release);
+        return refresh.promise;
+    }
+    async _readRelayRootSnapshot(transport) {
+        const root = ( URI.parse(ROOT_STATE_URI));
+        this._subscriptionManager.beginSnapshotRefresh(root);
+        try {
+            const result = await this._dispatchRequest("subscribe", {
+                channel: ROOT_STATE_URI
+            }, {
+                bypassInitializeQueue: true,
+                bypassReconnectGate: true,
+                timeoutMs: 30_000
+            });
+            if (transport !== this._transport) {
+                throw transportLostError(this._address);
+            }
+            const snapshot = result.snapshot;
+            if (!snapshot || !isAhpRootChannel(snapshot.resource) || !hasKey(snapshot.state, {
+                agents: true
+            })) {
+                throw ( new ProtocolError(
+                    JSON_RPC_INTERNAL_ERROR,
+                    "Relay root subscription returned no root snapshot"
+                ));
+            }
+            this._subscriptionManager.handleRootSnapshot(snapshot.state, snapshot.fromSeq);
+            this._serverSeq = Math.max(this._serverSeq, snapshot.fromSeq);
+        } finally {
+            if (transport === this._transport) {
+                this._subscriptionManager.cancelSnapshotRefresh(root);
+            }
+        }
+    }
     _normalizeAuthenticationParams(params) {
         return {
             ...params,
@@ -1105,11 +1535,14 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         }
         return toAgentHostUri(resource, this._connectionAuthority);
     }
-    async collectDebugLogs(session, kind, chat) {
+    async collectDebugLogs(session, kind, chat, progressToken) {
         const result = await this._sendExtensionRequest(CollectAgentHostDebugLogsExtensionMethod, {
             session: session?.toString(),
             chat: chat?.toString(),
-            kind
+            kind,
+            ...(progressToken ? {
+                progressToken
+            } : {})
         });
         if (result.kind !== kind) {
             throw ( new ProtocolError(
@@ -1164,10 +1597,23 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
             entries: result.entries
         };
     }
+    async cancelDebugLogs(progressToken) {
+        if (!supportsAgentHostDebugLogsCancellation(this._initializeResult.get())) {
+            this._logService.debug("[ExportAgentHostDebugLogs] Host does not support collection cancellation");
+            return;
+        }
+        await this._dispatchRequest(CancelAgentHostDebugLogsExtensionMethod, {
+            progressToken
+        }, {
+            timeoutMs: 10_000
+        });
+    }
     async readDebugLogsChunk(resource, position) {
-        const result = await this._sendExtensionRequest(ReadAgentHostDebugLogsChunkExtensionMethod, {
+        const result = await this._dispatchRequest(ReadAgentHostDebugLogsChunkExtensionMethod, {
             resource: ( fromAgentHostUri(resource).toString()),
             position
+        }, {
+            timeoutMs: 60_000
         });
         const data = decodeBase64(result.data);
         if (data.byteLength > AGENT_HOST_DEBUG_LOGS_CHUNK_BYTES) {
@@ -1195,6 +1641,9 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         await this._sendRequest("createChat", {
             channel: ( session.toString()),
             chat: ( chat.toString()),
+            ...(options?.workingDirectories !== undefined && !options.fork ? {
+                workingDirectories: options.workingDirectories.map(directory => ( fromAgentHostUri(directory).toString()))
+            } : {}),
             ...(options?.fork ? {
                 source: {
                     kind: ChatSourceKind.Fork,
@@ -1240,26 +1689,55 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         const result = await this._sendRequest("listSessions", {
             channel: ROOT_STATE_URI
         });
-        return result.items.map(s => ({
-            session: ( URI.parse(s.resource)),
-            startTime: Date.parse(s.createdAt),
-            modifiedTime: Date.parse(s.modifiedAt),
-            ...(s.project ? {
-                project: {
-                    uri: this._toClientUri(( URI.parse(s.project.uri))),
-                    displayName: s.project.displayName
-                }
-            } : {}),
-            summary: s.title,
-            status: s.status,
-            activity: s.activity,
-            workingDirectory: typeof s.workingDirectories?.[0] === "string" ? this._toClientUri(( URI.parse(s.workingDirectories[0]))) : undefined,
-            workingDirectories: s.workingDirectories?.map(d => this._toClientUri(( URI.parse(d)))),
-            changes: s.changes,
-            ...(s._meta !== undefined ? {
-                _meta: s._meta
-            } : {})
-        }));
+        return result.items.map(s => {
+            const model = readCodexSessionModel(s);
+            return {
+                session: ( URI.parse(s.resource)),
+                provider: s.provider,
+                startTime: Date.parse(s.createdAt),
+                modifiedTime: Date.parse(s.modifiedAt),
+                ...(s.project ? {
+                    project: {
+                        uri: this._toClientUri(( URI.parse(s.project.uri))),
+                        displayName: s.project.displayName
+                    }
+                } : {}),
+                summary: s.title,
+                status: s.status,
+                activity: s.activity,
+                workingDirectory: typeof s.workingDirectories?.[0] === "string" ? this._toClientUri(( URI.parse(s.workingDirectories[0]))) : undefined,
+                workingDirectories: s.workingDirectories?.map(d => this._toClientUri(( URI.parse(d)))),
+                changes: s.changes,
+                ...(model ? {
+                    model
+                } : {}),
+                chats: s.chats?.map(chat => ({
+                    chat: ( URI.parse(chat.resource)),
+                    summary: chat.title,
+                    kind: s.defaultChat === chat.resource || isDefaultChatUri(chat.resource) ? "default" : "peer",
+                    origin: chat.origin,
+                    ...(chat.interactivity !== undefined ? {
+                        interactivity: chat.interactivity
+                    } : {}),
+                    ...(isSessionChatArchived(chat) ? {
+                        archived: true
+                    } : {}),
+                    ...(chat.status !== undefined ? {
+                        status: chat.status,
+                        isRead: isSessionStatusRead(chat.status)
+                    } : {}),
+                    ...(chat.changes !== undefined ? {
+                        changes: chat.changes
+                    } : {})
+                })) ?? (s.defaultChat ? [{
+                    chat: ( URI.parse(s.defaultChat)),
+                    kind: "default"
+                }] : undefined),
+                ...(s._meta !== undefined ? {
+                    _meta: s._meta
+                } : {})
+            };
+        });
     }
     _toClientUri(uri) {
         return uri.scheme === Schemas.file ? toAgentHostUri(uri, this._connectionAuthority) : uri;
@@ -1358,8 +1836,7 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         if (this._state.kind === AgentHostClientState.Closed) {
             return;
         }
-        this._lastReadTime = Date.now();
-        this._resetLivenessTimers();
+        this._recordInboundTraffic();
         if (isJsonRpcRequest(msg)) {
             this._handleReverseRequest(msg.id, msg.method, msg.params);
         } else if (isJsonRpcResponse(msg)) {
@@ -1382,6 +1859,9 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
                 );
             }
         } else if (isJsonRpcNotification(msg)) {
+            if (this._devContainerService.handleNotification(msg.method, msg.params)) {
+                return;
+            }
             switch (msg.method) {
             case "action":
                 {
@@ -1433,6 +1913,7 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         }
     }
     _handleFatalClose(error, reason) {
+        this._relayKeepAliveTimer.cancel();
         this._onDidFatalClose.fire(error);
         this._handleClose(error, reason);
     }
@@ -1440,6 +1921,15 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         if (this._state.kind === AgentHostClientState.Closed) {
             return;
         }
+        this._onDidConnectionDiagnostic.fire({
+            operationId: this._clientId,
+            attemptId: this._diagnosticAttemptId,
+            phase: "connection.closed",
+            outcome: "info",
+            timestamp: Date.now(),
+            error: getConnectionDiagnosticError(error),
+            detail: reason
+        });
         this._cancelLivenessTimers();
         if (this._state.kind === AgentHostClientState.Reconnecting) {
             const reconnect = this._state.reconnect;
@@ -1469,7 +1959,7 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         }
         let closeListener = Disposable.None;
         const closePromise = ( new Promise((_resolve, reject) => {
-            closeListener = this.onDidClose(() => reject(
+            closeListener = Event.once(this.onDidClose)(() => reject(
                 this._state.kind === AgentHostClientState.Closed ? this._state.error : connectionClosedError(this._address)
             ));
         }));
@@ -1546,6 +2036,23 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         void (async () => {
             try {
                 switch (method) {
+                case RequestAgentHostMcpAuthenticationExtensionMethod:
+                    {
+                        const request = mcpAuthenticationRequestValidator.validateOrThrow(params);
+                        let authenticated = false;
+                        try {
+                            authenticated = (await this._mcpAuthenticationHandler?.(request)) ?? false;
+                        } catch (err) {
+                            this._logService.error(
+                                "[AgentHostProtocolClient] Failed to silently authenticate MCP server",
+                                err
+                            );
+                        }
+                        sendResult({
+                            authenticated
+                        });
+                        return;
+                    }
                 case RequestAgentHostWorkspaceTrustExtensionMethod:
                     {
                         if (typeof p.workspace !== "string") {
@@ -1583,7 +2090,7 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
                         const trusted = await this._workspaceTrustRequestService.requestResourcesTrust({
                             uri: workspace,
                             message: ( localize(
-                                1848,
+                                1863,
                                 "An agent session will be able to read files, run commands, and make changes in this folder."
                             ))
                         });
@@ -1722,6 +2229,9 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
     _sendExtensionRequest(method, params) {
         return this._dispatchRequest(method, params);
     }
+    sendHostExtensionRequest(method, params) {
+        return this._dispatchRequest(method, params);
+    }
     _updateTelemetryLevel() {
         this._dispatchRootConfig({
             [AgentHostTelemetryLevelConfigKey]: telemetryLevelToAgentHostConfigValue(this._effectiveTelemetryLevel())
@@ -1734,6 +2244,9 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         );
     }
     _dispatchRootConfig(config) {
+        if (this._isWebPubSubRelay()) {
+            return;
+        }
         this.dispatchAction(ROOT_STATE_URI, {
             type: ActionType.RootConfigChanged,
             config
@@ -1757,6 +2270,9 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         });
     }
     _updateManagedSettingsPermissions(sendDuringReconnect = false) {
+        if (this._isWebPubSubRelay()) {
+            return;
+        }
         const permissions = this._resourceIdentity === LOCAL_AGENT_HOST_RESOURCE_IDENTITY ? resolveManagedSettingsPermissions(this._configurationService) : {};
         this._sendExtensionNotification("setClientManagedSettingsPermissions", {
             permissions
@@ -1773,7 +2289,7 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
             const {
                 request,
                 result
-            } = this._createRequest(method, params);
+            } = this._createRequest(method, params, options.timeoutMs);
             this._transport.send(request);
             return result;
         }
@@ -1781,7 +2297,7 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
             const {
                 request,
                 result
-            } = this._createRequest(method, params);
+            } = this._createRequest(method, params, options.timeoutMs);
             this._state.outbox.push(request);
             return result;
         }
@@ -1801,11 +2317,23 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         const {
             request,
             result
-        } = this._createRequest(method, params);
+        } = this._createRequest(method, params, options.timeoutMs);
+        if (method === "initialize") {
+            this._logService.info(
+                `[RemoteAgentHostProtocol] Sending initialize: address=${this._address} clientId=${this._clientId} requestId=${request.id}`
+            );
+        }
+        if (this._firstSessionRequestPending && current.kind === AgentHostClientState.Connected && (method === "createSession" || method === "listSessions" || method === "subscribe")) {
+            this._firstSessionRequestPending = false;
+            return this._traceConnection("protocol.firstSessionRequest", async () => {
+                this._transport.send(request);
+                return result;
+            });
+        }
         this._transport.send(request);
         return result;
     }
-    _createRequest(method, params) {
+    _createRequest(method, params, timeoutMs) {
         const id = this._nextRequestId++;
         const deferred = ( new DeferredPromise());
         this._pendingRequests.set(id, {
@@ -1813,6 +2341,15 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
             suppressNotFoundWarning: isFileResourceRead(method, params),
             sentAt: Date.now()
         });
+        if (timeoutMs !== undefined) {
+            const deadline = disposableTimeout(() => {
+                if (this._pendingRequests.get(id)?.deferred === deferred) {
+                    this._pendingRequests.delete(id);
+                    deferred.error(( new ProtocolError(JSON_RPC_INTERNAL_ERROR, `${method} response timed out`)));
+                }
+            }, timeoutMs, this._store);
+            void deferred.p.then(() => deadline.dispose(), () => deadline.dispose());
+        }
         return {
             request: {
                 jsonrpc: "2.0",
@@ -1838,28 +2375,70 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         }
         this._pendingRequests.clear();
     }
+    _recordInboundTraffic() {
+        this._lastReadTime = Date.now();
+        this._resetLivenessTimers();
+    }
+    _startRelayKeepAlive(meta) {
+        this._relayKeepAliveTimer.cancel();
+        const window = readRelayKeepAliveTimeout({
+            _meta: meta
+        });
+        if (!this._isWebPubSubRelay() || window === undefined) {
+            return;
+        }
+        this._relayKeepAliveTimer.cancelAndSet(() => {
+            if (this._state.kind !== AgentHostClientState.Connected || this._relayKeepAlivePending) {
+                return;
+            }
+            const operation = this._dispatchRequest("ping", {
+                channel: ROOT_STATE_URI
+            }, {
+                timeoutMs: 10_000
+            });
+            this._relayKeepAlivePending = operation;
+            void operation.catch(error => {
+                this._logService.warn("[AgentHostProtocolClient] Relay keep-alive failed", error);
+            }).finally(() => {
+                if (this._relayKeepAlivePending === operation) {
+                    this._relayKeepAlivePending = undefined;
+                }
+            });
+        }, Math.max(1, Math.min(60_000, Math.floor(window / 3))));
+    }
     _resetLivenessTimers() {
         this._cancelLivenessTimers();
-        if (this._state.kind === AgentHostClientState.Incompatible || this._state.kind === AgentHostClientState.Reconnecting || this._state.kind === AgentHostClientState.Closed) {
+        if (!this._canCheckLiveness()) {
             return;
         }
         this._pingTimer.cancelAndSet(() => this._onPingTimer(), PING_INTERVAL_MS);
         this._closeTimer.cancelAndSet(() => this._onCloseTimer(), PING_INTERVAL_MS + LIVENESS_TIMEOUT_MS);
     }
+    _isWebPubSubRelay() {
+        return this._transport.clientConnectionKind === AgentHostClientConnectionKind.WebPubSub || this._transport.clientConnectionKind === AgentHostClientConnectionKind.MissionControl;
+    }
     _cancelLivenessTimers() {
         this._pingTimer.cancel();
         this._closeTimer.cancel();
         this._livenessDeferred = false;
         this._livenessDeferredSince = undefined;
     }
+    _canCheckLiveness() {
+        return this._state.kind !== AgentHostClientState.Incompatible && this._state.kind !== AgentHostClientState.Closed && (this._state.kind !== AgentHostClientState.Reconnecting || this._state.reconnect.transportConnected);
+    }
     _onPingTimer() {
-        if (this._state.kind === AgentHostClientState.Incompatible || this._state.kind === AgentHostClientState.Closed || this._state.kind === AgentHostClientState.Reconnecting) {
+        if (!this._canCheckLiveness()) {
             return;
         }
-        void this.ping().catch(() => undefined);
+        void this._dispatchRequest("ping", {
+            channel: ROOT_STATE_URI
+        }, {
+            bypassReconnectGate: true,
+            bypassInitializeQueue: true
+        }).catch(() => undefined);
     }
     _onCloseTimer() {
-        if (this._state.kind === AgentHostClientState.Incompatible || this._state.kind === AgentHostClientState.Closed || this._state.kind === AgentHostClientState.Reconnecting) {
+        if (!this._canCheckLiveness()) {
             return;
         }
         if (this._transport.clientConnectionKind === AgentHostClientConnectionKind.Local) {
@@ -1879,8 +2458,9 @@ let AgentHostProtocolClient = class AgentHostProtocolClient extends Disposable {
         }
         const silence = Date.now() - this._lastReadTime;
         this._logService.info(
-            `[RemoteAgentHostProtocol] Liveness: no message from ${this._address} for ${silence}ms; forcing close to trigger reconnect.`
+            `[RemoteAgentHostProtocol] Liveness: no data from ${this._address} for ${silence}ms; forcing close to trigger reconnect.`
         );
+        this._diagnostic("watchdog.timeout", `sinceLastDataMs=${silence}`);
         this._transportListeners.clear();
         if (this._transportFactory) {
             this._rejectPendingRequests(connectionTimeoutError(this._address, silence));