automations: fix: use aggregate discovery for New badge

Include legacy ledger readability and live Agent Host catalogue health in badge eligibility. Propagate host disconnects, retain simple prior-window tracking, and add regression coverage without changing shared storage APIs.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 7e75f732-f18f-42a7-b1e2-a034c03596eb

Author: ulugbekna

src/vs/sessions/AUTOMATIONS.md

@@ -216,16 +216,6 @@ Every non-supported state exposes the provider-scoped legacy fallback when it co
 
 After legacy data has migrated and its source is drained, disconnecting or disabling the provider does not create a second authority. The host retains durable definitions; its projection becomes available again after reconnect or re-enable.
 
-### Initial discovery
-
-Provider Automation stores may expose an initial-discovery state through `ISessionsProviderAutomations`:
-
-- `pending` means the provider may still reveal existing definitions;
-- `ready` means its initial authoritative catalogue is definitive;
-- `unavailable` means it cannot currently establish a definitive catalogue.
-
-Stores without this state are synchronously initialized. Agent Host stores become ready only after their initial catalogue and migration state are authoritative. A host without Automation capability is ready because its synchronous legacy store remains authoritative; a disconnected host is unavailable because host-owned definitions cannot be projected.
-
 ## Migration
 
 Migration has two ownership boundaries:

src/vs/sessions/contrib/automations/test/browser/providerAutomationService.test.ts

@@ -1257,13 +1257,16 @@ suite('ProviderAutomationService', () => {
 			runs: [],
 		});
 		const { service, providerStore, storage } = createService(futureLedger);
+		const initialDiscoveryState = service.initialDiscoveryState.get();
 
 		await assert.rejects(service.waitForMigrationForTesting(), /cannot be migrated safely/);
 
 		assert.deepStrictEqual({
+			initialDiscoveryState,
 			providerAutomations: providerStore.automations.get(),
 			persisted: storage.get(AUTOMATION_STORAGE_KEY, StorageScope.APPLICATION),
 		}, {
+			initialDiscoveryState: 'unavailable',
 			providerAutomations: [],
 			persisted: futureLedger,
 		});

src/vs/sessions/contrib/providers/agentHost/browser/agentHostAutomationStore.ts

@@ -26,7 +26,7 @@ import { ITelemetryService } from '../../../../../platform/telemetry/common/tele
 import { assertAutomationSessionTemplate, type AutomationRunTrigger, type AutomationTarget, type IAutomationDescriptor, type IAutomationRun, type IAutomationSchedule, type IAutomationSessionTemplate } from '../../../../../workbench/contrib/chat/common/automations/automation.js';
 import { AutomationActiveRunError, type AutomationCatalogueState, assertAutomationSessionTemplateAuthority, combineAutomationCatalogueStates, type AutomationMutationGuard, type IAutomationRunClaim, type ICreateAutomationOptions, type IGuardedAutomationUpdateResult, isAutomationActiveRunError, serializeAutomationEditableState, type IUpdateAutomationOptions, type IUpdateAutomationRunOptions } from '../../../../../workbench/contrib/chat/common/automations/automationService.js';
 import { publishAutomationMigration } from '../../../../../workbench/contrib/chat/common/automations/automationTelemetry.js';
-import type { AutomationInitialDiscoveryState, IAutomation, IAutomationSnapshotImportResult, IGuardedAutomationSnapshotRemovalResult, ISessionsProviderAutomations } from '../../../../services/sessions/common/sessionsProvider.js';
+import type { IAutomation, IAutomationSnapshotImportResult, IGuardedAutomationSnapshotRemovalResult, ISessionsProviderAutomations } from '../../../../services/sessions/common/sessionsProvider.js';
 import { IAutomationStorageService } from '../../../automations/common/automationStorageService.js';
 
 const MUTATION_TIMEOUT_MS = 30_000;
@@ -87,7 +87,6 @@ export class AgentHostAutomationStore extends Disposable implements ISessionsPro
 	private _migrationPromise: Promise<void> | undefined;
 	private _lastPreflightDeferralKey: string | undefined;
 
-	readonly initialDiscoveryState = derived<AutomationInitialDiscoveryState>(this, reader => this._ready.read(reader) ? 'ready' : 'pending');
 	readonly automations: IObservable<readonly IAutomationDescriptor[]>;
 	readonly runs: IObservable<readonly IAutomationRun[]>;
 	readonly catalogueState: IObservable<AutomationCatalogueState>;

src/vs/sessions/contrib/providers/agentHost/browser/localAgentHostSessionsProvider.ts

@@ -264,6 +264,10 @@ export class LocalAgentHostSessionsProvider extends BaseAgentHostSessionsProvide
 		};
 		bindConnection();
 		this._register(this._agentHostService.onAgentHostStart(bindConnection));
+		this._register(this._agentHostService.onAgentHostExit(() => {
+			connectionListeners.clear();
+			automations.clearConnection();
+		}));
 
 		// Eagerly populate the session cache once authentication has settled.
 		// Without this, the sidebar would only call `getSessions()` after some

src/vs/sessions/contrib/providers/agentHost/browser/reconnectableAgentHostAutomationStore.ts

@@ -13,7 +13,7 @@ import { IConfigurationService } from '../../../../../platform/configuration/com
 import { ILogService } from '../../../../../platform/log/common/log.js';
 import type { AutomationRunTrigger, IAutomationDescriptor, IAutomationRun } from '../../../../../workbench/contrib/chat/common/automations/automation.js';
 import { type AutomationCatalogueState, isAutomationActiveRunError, type AutomationMutationGuard, type IAutomationRunClaim, type ICreateAutomationOptions, type IGuardedAutomationUpdateResult, type IUpdateAutomationOptions, type IUpdateAutomationRunOptions } from '../../../../../workbench/contrib/chat/common/automations/automationService.js';
-import type { AutomationInitialDiscoveryState, IAutomation, IAutomationSnapshotImportResult, IGuardedAutomationSnapshotRemovalResult, ISessionsProviderAutomations } from '../../../../services/sessions/common/sessionsProvider.js';
+import type { IAutomation, IAutomationSnapshotImportResult, IGuardedAutomationSnapshotRemovalResult, ISessionsProviderAutomations } from '../../../../services/sessions/common/sessionsProvider.js';
 import { AgentHostAutomationStore, type IAgentHostAutomationBoundaryMapper, type IAgentHostAutomationConnection } from './agentHostAutomationStore.js';
 import { CHAT_AUTOMATIONS_ENABLED_SETTING } from '../../../../../workbench/contrib/chat/common/automations/automationsEnabled.js';
 
@@ -35,20 +35,6 @@ export class ReconnectableAgentHostAutomationStore extends Disposable implements
 	private readonly _authorityState = observableValue<AutomationAuthorityState>(this, { kind: 'disconnected' });
 	private readonly _disposeCancellation = new CancellationTokenSource();
 
-	readonly initialDiscoveryState = derived<AutomationInitialDiscoveryState>(this, reader => {
-		const state = this._authorityState.read(reader);
-		switch (state.kind) {
-			case 'supported':
-				return state.store.initialDiscoveryState.read(reader);
-			case 'unsupported':
-				return 'ready';
-			case 'initializing':
-				return 'pending';
-			case 'disconnected':
-			case 'disabled':
-				return 'unavailable';
-		}
-	});
 	readonly automations = derived(this, reader => this._currentStore.read(reader)?.automations.read(reader) ?? this._legacySource?.automations.read(reader) ?? []);
 	readonly runs = derived(this, reader => this._currentStore.read(reader)?.runs.read(reader) ?? this._legacySource?.runs.read(reader) ?? []);
 	readonly catalogueState: IObservable<AutomationCatalogueState> = derived(this, reader => {

src/vs/sessions/contrib/providers/agentHost/test/browser/agentHostAutomationStore.test.ts

@@ -2214,19 +2214,16 @@ suite('AgentHostAutomationStore', () => {
 			configurationService,
 		));
 		store.setConnection(connection);
-		assert.strictEqual(store.initialDiscoveryState.get(), 'pending');
 		await assert.rejects(store.completeMigration(), /cannot be migrated safely/);
 
 		legacy.migrationAllowed = true;
 		store.setConnection(connection);
 		await store.completeMigration();
 
 		assert.deepStrictEqual({
-			initialDiscoveryState: store.initialDiscoveryState.get(),
 			subscriptions: connection.subscribedChannel,
 			completionRequests: connection.dispatched.filter(entry => entry.channel === ROOT_STATE_URI).length,
 		}, {
-			initialDiscoveryState: 'ready',
 			subscriptions: URI.parse(AUTOMATION_CATALOG_URI).toString(),
 			completionRequests: 1,
 		});
@@ -2416,7 +2413,6 @@ suite('AgentHostAutomationStore', () => {
 			configurationService,
 		));
 		store.setConnection(connection);
-		const pendingState = store.initialDiscoveryState.get();
 
 		let settled = false;
 		const migration = store.completeMigration().finally(() => settled = true);
@@ -2430,15 +2426,7 @@ suite('AgentHostAutomationStore', () => {
 		}, undefined);
 		await migration;
 
-		assert.deepStrictEqual({
-			pendingState,
-			readyState: store.initialDiscoveryState.get(),
-			completionRequests: connection.dispatched.filter(entry => entry.channel === ROOT_STATE_URI).length,
-		}, {
-			pendingState: 'pending',
-			readyState: 'ready',
-			completionRequests: 1,
-		});
+		assert.strictEqual(connection.dispatched.filter(entry => entry.channel === ROOT_STATE_URI).length, 1);
 	});
 
 	test('migration resolves without subscribing after an older host finishes initializing', async () => {
@@ -2463,7 +2451,6 @@ suite('AgentHostAutomationStore', () => {
 		));
 		store.setConnection(connection);
 		const migration = store.completeMigration();
-		const pendingState = store.initialDiscoveryState.get();
 
 		connection.initializeResult.set({
 			protocolVersion: '1',
@@ -2473,13 +2460,9 @@ suite('AgentHostAutomationStore', () => {
 		await migration;
 
 		assert.deepStrictEqual({
-			pendingState,
-			readyState: store.initialDiscoveryState.get(),
 			subscribedChannel: connection.subscribedChannel,
 			completionRequests: connection.dispatched.filter(entry => entry.channel === ROOT_STATE_URI).length,
 		}, {
-			pendingState: 'pending',
-			readyState: 'ready',
 			subscribedChannel: undefined,
 			completionRequests: 0,
 		});
@@ -2509,13 +2492,7 @@ suite('AgentHostAutomationStore', () => {
 
 		await store.completeMigration();
 
-		assert.deepStrictEqual({
-			initialDiscoveryState: store.initialDiscoveryState.get(),
-			dispatched: connection.dispatched.length,
-		}, {
-			initialDiscoveryState: 'pending',
-			dispatched: 0,
-		});
+		assert.strictEqual(connection.dispatched.length, 0);
 	}));
 
 	test('disposing while capabilities initialize settles migration', async () => {

src/vs/sessions/contrib/providers/agentHost/test/browser/localAgentHostSessionsProvider.test.ts

@@ -38,6 +38,7 @@ import { IWorkspaceTrustManagementService, IWorkspaceTrustRequestService, Resour
 import { IChatWidget, IChatWidgetService } from '../../../../../../workbench/contrib/chat/browser/chat.js';
 import { IChatService, type ChatSendResult, type IChatModelReference, type IChatSendRequestOptions } from '../../../../../../workbench/contrib/chat/common/chatService/chatService.js';
 import { IChatSessionsService, isIChatSessionFileChange2 } from '../../../../../../workbench/contrib/chat/common/chatSessionsService.js';
+import { CHAT_AUTOMATIONS_ENABLED_SETTING } from '../../../../../../workbench/contrib/chat/common/automations/automationsEnabled.js';
 import { ChatModeKind } from '../../../../../../workbench/contrib/chat/common/constants.js';
 import { ILanguageModelsService, type ILanguageModelChatMetadata } from '../../../../../../workbench/contrib/chat/common/languageModels.js';
 import type { IChatModel, IChatModelInputState, IInputModel } from '../../../../../../workbench/contrib/chat/common/model/chatModel.js';
@@ -88,6 +89,8 @@ class MockAgentHostService extends mock<IAgentHostService>() {
 	override get rootState(): IAgentSubscription<RootState> { return this._rootStateSubscription; }
 	private readonly _onAgentHostStart = new Emitter<void>();
 	override readonly onAgentHostStart = this._onAgentHostStart.event;
+	private readonly _onAgentHostExit = new Emitter<number>();
+	override readonly onAgentHostExit = this._onAgentHostExit.event;
 	override readonly initializeResult = constObservable({
 		protocolVersion: '1',
 		serverSeq: 0,
@@ -381,6 +384,10 @@ class MockAgentHostService extends mock<IAgentHostService>() {
 		this._onAgentHostStart.fire();
 	}
 
+	fireAgentHostExit(): void {
+		this._onAgentHostExit.fire(0);
+	}
+
 	setRootStateError(): void {
 		const error = new Error('root state failed');
 		this._rootStateValue = error;
@@ -401,6 +408,7 @@ class MockAgentHostService extends mock<IAgentHostService>() {
 		this._onDidRootStateChange.dispose();
 		this._onDidRootStateError.dispose();
 		this._onAgentHostStart.dispose();
+		this._onAgentHostExit.dispose();
 		for (const emitter of this._sessionStateEmitters.values()) {
 			emitter.dispose();
 		}
@@ -672,6 +680,27 @@ suite('LocalAgentHostSessionsProvider', () => {
 
 	// ---- Provider identity -------
 
+	test('Automation discovery follows local Agent Host connection lifetime', () => {
+		const provider = createProvider(disposables, agentHost, undefined, {
+			configurationService: new TestConfigurationService({ [CHAT_AUTOMATIONS_ENABLED_SETTING]: true }),
+		});
+		const initial = provider.automations.initialDiscoveryState?.get();
+
+		agentHost.fireAgentHostExit();
+		const disconnected = provider.automations.initialDiscoveryState?.get();
+		agentHost.fireAgentHostStart();
+
+		assert.deepStrictEqual({
+			initial,
+			disconnected,
+			reconnected: provider.automations.initialDiscoveryState?.get(),
+		}, {
+			initial: 'ready',
+			disconnected: 'unavailable',
+			reconnected: 'ready',
+		});
+	});
+
 	test('has correct id, label, and sessionType from rootState agents', () => {
 		const provider = createProvider(disposables, agentHost);
 

src/vs/sessions/contrib/sessions/browser/automationsNewBadge.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { Disposable, MutableDisposable } from '../../../../base/common/lifecycle.js';
-import { autorun, derived, IObservable, observableSignalFromEvent, observableValue } from '../../../../base/common/observable.js';
+import { autorun, derived, observableValue } from '../../../../base/common/observable.js';
 import { onUnexpectedError } from '../../../../base/common/errors.js';
 import { IConfigurationService, isConfigured } from '../../../../platform/configuration/common/configuration.js';
 import { ILogService } from '../../../../platform/log/common/log.js';
@@ -14,7 +14,6 @@ import { IAutomationService } from '../../../../workbench/contrib/chat/common/au
 import { IWorkbenchAssignmentService } from '../../../../workbench/services/assignment/common/assignmentService.js';
 import { ILifecycleService, LifecyclePhase } from '../../../../workbench/services/lifecycle/common/lifecycle.js';
 import { ICustomViewService } from '../../../services/customView/browser/customViewService.js';
-import { ISessionsProvidersService } from '../../../services/sessions/browser/sessionsProvidersService.js';
 import { ISessionsWindowUsageService } from '../../../services/sessions/browser/sessionsWindowUsageService.js';
 import { AUTOMATIONS_CUSTOM_VIEW_ID } from './automationsConstants.js';
 
@@ -43,7 +42,6 @@ export class AutomationsNewBadgeState extends Disposable {
 	private readonly forcePreview = observableValue(this, false);
 	private readonly startupDecision = observableValue<AutomationsNewBadgeStartupDecision>(this, 'pending');
 	private readonly automationEvidenceObserver = this._register(new MutableDisposable());
-	private readonly providersChanged: IObservable<void>;
 	private observersRegistered = false;
 	private initializationPromise: Promise<void> | undefined;
 	private styleRequest = 0;
@@ -66,12 +64,10 @@ export class AutomationsNewBadgeState extends Disposable {
 		@IConfigurationService private readonly configurationService: IConfigurationService,
 		@ILogService private readonly logService: ILogService,
 		@ISessionsWindowUsageService private readonly sessionsWindowUsageService: ISessionsWindowUsageService,
-		@ISessionsProvidersService private readonly sessionsProvidersService: ISessionsProvidersService,
 		@ILifecycleService private readonly lifecycleService: ILifecycleService,
 	) {
 		super();
 		this.seen = this._register(automationsNewBadgeSeenMemento(StorageScope.APPLICATION, StorageTarget.MACHINE, storageService));
-		this.providersChanged = observableSignalFromEvent(this, sessionsProvidersService.onDidChangeProviders);
 	}
 
 	initialize(): Promise<void> {
@@ -100,13 +96,8 @@ export class AutomationsNewBadgeState extends Disposable {
 				if (this.forcePreview.read(reader) || this.seen.read(reader) || this.startupDecision.read(reader) !== 'eligible') {
 					return;
 				}
-				this.providersChanged.read(reader);
-				for (const provider of this.sessionsProvidersService.getProviders()) {
-					if (provider.automations
-						&& (provider.automations.initialDiscoveryState?.read(reader) ?? 'ready') !== 'ready') {
-						this.startupDecision.set('suppressed', undefined);
-						return;
-					}
+				if (this.automationService.initialDiscoveryState.read(reader) !== 'ready') {
+					this.startupDecision.set('suppressed', undefined);
 				}
 			}));
 		}
@@ -133,7 +124,7 @@ export class AutomationsNewBadgeState extends Disposable {
 	}
 
 	private async doInitialize(): Promise<void> {
-		if (this.seen.get() || this.forcePreview.get()) {
+		if (this._store.isDisposed || this.seen.get() || this.forcePreview.get()) {
 			return;
 		}
 
@@ -143,14 +134,11 @@ export class AutomationsNewBadgeState extends Disposable {
 		}
 
 		await this.lifecycleService.when(LifecyclePhase.Eventually);
-		if (this.seen.get() || this.forcePreview.get()) {
+		if (this._store.isDisposed || this.seen.get() || this.forcePreview.get()) {
 			return;
 		}
 
-		const providersReady = this.sessionsProvidersService.getProviders()
-			.every(provider => !provider.automations
-				|| (provider.automations.initialDiscoveryState?.get() ?? 'ready') === 'ready');
-		if (!providersReady) {
+		if (this.automationService.initialDiscoveryState.get() !== 'ready') {
 			this.startupDecision.set('suppressed', undefined);
 			return;
 		}
@@ -187,7 +175,7 @@ export class AutomationsNewBadgeState extends Disposable {
 	}
 
 	private canResolveStyle(): boolean {
-		return !this.seen.get() && (this.forcePreview.get() || this.startupDecision.get() === 'eligible');
+		return !this._store.isDisposed && !this.seen.get() && (this.forcePreview.get() || this.startupDecision.get() === 'eligible');
 	}
 
 	private normalizeStyle(value: string | undefined): AutomationsNewBadgeStyle {