sessions: Preserve custom group for starting sessions

src/vs/sessions/contrib/sessions/test/browser/sessionsArchiveActions.test.ts

@@ -58,9 +58,11 @@ suite('Sessions - Bulk archive undo', () => {
 			override readonly onDidDeleteSession = Event.None;
 			override readonly onWillSendRequest = Event.None;
 			override readonly onDidReplaceSession = Event.None;
+			override readonly onDidReplaceNewDraftSession = Event.None;
 			override readonly onDidStartSession = Event.None;
 			override readonly onDidDiscardNewSession = Event.None;
 			override getSessions() { return sessions.map(entry => entry.session); }
+			override getInFlightNewSessionRequests() { return []; }
 			override getSession(resource: URI) {
 				return sessions.find(entry => entry.session.resource.toString() === resource.toString())?.session;
 			}

src/vs/sessions/services/sessions/browser/sessionGroupsService.ts

@@ -158,15 +158,34 @@ export class SessionGroupsService extends Disposable implements ISessionGroupsSe
 		// membership here would turn that transient gap into a permanent,
 		// unrecoverable loss of the user's grouping.
 		this._register(this.sessionsManagementService.onDidChangeSessions(e => {
+			const transientMembershipChanged = new Set<string>();
+			const inFlightSessionIds = new Set(this.sessionsManagementService.getInFlightNewSessionRequests().map(session => session.sessionId));
+			for (const session of e.added) {
+				if (inFlightSessionIds.has(session.sessionId) && this.capturePendingNewSessionGroup(session.sessionId)) {
+					transientMembershipChanged.add(session.sessionId);
+				}
+			}
 			for (const session of e.removed) {
-				this._inFlightSessionGroups.delete(session.sessionId);
+				if (this._inFlightSessionGroups.delete(session.sessionId)) {
+					transientMembershipChanged.add(session.sessionId);
+				}
+			}
+			const visibleSessionIds = new Set(this.sessionsManagementService.getSessions().map(session => session.sessionId));
+			for (const sessionId of this._inFlightSessionGroups.keys()) {
+				if (!visibleSessionIds.has(sessionId)) {
+					this._inFlightSessionGroups.delete(sessionId);
+					transientMembershipChanged.add(sessionId);
+				}
 			}
 			const changed = new Set<string>();
 			const archivedStateChanged = this.removeArchivedMembership([...e.added, ...e.changed], changed);
 			this.updateDefaultPlacement(this.sessionsManagementService.getSessions(), changed);
 			if (archivedStateChanged || changed.size > 0) {
 				this.save();
 			}
+			for (const sessionId of transientMembershipChanged) {
+				changed.add(sessionId);
+			}
 			if (changed.size > 0) {
 				this._onDidChange.fire({ groupsChanged: false, membershipChanged: changed });
 			}
@@ -199,24 +218,14 @@ export class SessionGroupsService extends Disposable implements ISessionGroupsSe
 		// can no longer rebind it. A send into an existing session discards the
 		// draft (firing the discard handler below) before this fires.
 		this._register(this.sessionsManagementService.onWillSendRequest(session => {
-			if (this._pendingNewSessionGroupId === undefined) {
-				return;
+			if (this.capturePendingNewSessionGroup(session.sessionId)) {
+				this._onDidChange.fire({ groupsChanged: false, membershipChanged: new Set([session.sessionId]) });
 			}
-			this._inFlightSessionGroups.set(session.sessionId, this._pendingNewSessionGroupId);
-			this._pendingNewSessionGroupId = undefined;
 		}));
 
-		// A draft graduates into a committed session with a new id; follow it.
-		this._register(this.sessionsManagementService.onDidReplaceSession(({ from, to }) => {
-			if (from.sessionId === to.sessionId) {
-				return;
-			}
-			const groupId = this._inFlightSessionGroups.get(from.sessionId);
-			if (groupId !== undefined) {
-				this._inFlightSessionGroups.delete(from.sessionId);
-				this._inFlightSessionGroups.set(to.sessionId, groupId);
-			}
-		}));
+		// Follow both pre-send draft replacement and committed-session graduation.
+		this._register(this.sessionsManagementService.onDidReplaceNewDraftSession(({ from, to }) => this.replaceInFlightSession(from, to)));
+		this._register(this.sessionsManagementService.onDidReplaceSession(({ from, to }) => this.replaceInFlightSession(from, to)));
 
 		// The started session carries the committed id; record its group now.
 		this._register(this.sessionsManagementService.onDidStartSession(session => {
@@ -243,7 +252,7 @@ export class SessionGroupsService extends Disposable implements ISessionGroupsSe
 		do {
 			placed = false;
 			for (const session of sessions) {
-				if (session.isArchived.get() || this._membership.has(session.sessionId) || this._explicitlyUngroupedSessionIds.has(session.sessionId)) {
+				if (session.isArchived.get() || this._membership.has(session.sessionId) || this._inFlightSessionGroups.has(session.sessionId) || this._explicitlyUngroupedSessionIds.has(session.sessionId)) {
 					continue;
 				}
 				const creatorResource = session.createdBySession?.get()?.session;
@@ -272,7 +281,7 @@ export class SessionGroupsService extends Disposable implements ISessionGroupsSe
 		const membershipChanged = new Set<string>();
 		if (memberSessionIds) {
 			for (const sessionId of memberSessionIds) {
-				this.setMembership(sessionId, group.id, membershipChanged);
+				this.setCurrentMembership(sessionId, group.id, membershipChanged);
 			}
 		}
 		this.updateDefaultPlacement(this.sessionsManagementService.getSessions(), membershipChanged);
@@ -323,7 +332,7 @@ export class SessionGroupsService extends Disposable implements ISessionGroupsSe
 		const sessionIds = typeof sessionIdOrIds === 'string' ? [sessionIdOrIds] : sessionIdOrIds;
 		const membershipChanged = new Set<string>();
 		for (const sessionId of sessionIds) {
-			this.setMembership(sessionId, groupId, membershipChanged);
+			this.setCurrentMembership(sessionId, groupId, membershipChanged);
 		}
 		this.updateDefaultPlacement(this.sessionsManagementService.getSessions(), membershipChanged);
 		if (membershipChanged.size === 0) {
@@ -334,7 +343,9 @@ export class SessionGroupsService extends Disposable implements ISessionGroupsSe
 	}
 
 	removeFromGroup(sessionId: string): void {
-		if (!this._membership.delete(sessionId)) {
+		const membershipDeleted = this._membership.delete(sessionId);
+		const transientMembershipDeleted = this._inFlightSessionGroups.delete(sessionId);
+		if (!membershipDeleted && !transientMembershipDeleted) {
 			return;
 		}
 		this.markExplicitlyUngrouped(sessionId);
@@ -343,17 +354,22 @@ export class SessionGroupsService extends Disposable implements ISessionGroupsSe
 	}
 
 	getGroupOfSession(sessionId: string): string | undefined {
-		return this._membership.get(sessionId);
+		return this._membership.get(sessionId) ?? this._inFlightSessionGroups.get(sessionId);
 	}
 
 	getSessionIdsInGroup(groupId: string): string[] {
-		const result: string[] = [];
+		const result = new Set<string>();
 		for (const [sessionId, gid] of this._membership) {
 			if (gid === groupId) {
-				result.push(sessionId);
+				result.add(sessionId);
 			}
 		}
-		return result;
+		for (const [sessionId, gid] of this._inFlightSessionGroups) {
+			if (gid === groupId) {
+				result.add(sessionId);
+			}
+		}
+		return [...result];
 	}
 
 	setPendingNewSessionGroup(groupId: string): void {
@@ -362,6 +378,40 @@ export class SessionGroupsService extends Disposable implements ISessionGroupsSe
 
 	// -- Helpers --
 
+	private capturePendingNewSessionGroup(sessionId: string): boolean {
+		if (this._pendingNewSessionGroupId === undefined) {
+			return false;
+		}
+		this._inFlightSessionGroups.set(sessionId, this._pendingNewSessionGroupId);
+		this._pendingNewSessionGroupId = undefined;
+		return true;
+	}
+
+	private replaceInFlightSession(from: ISession, to: ISession): void {
+		if (from.sessionId === to.sessionId) {
+			return;
+		}
+		const groupId = this._inFlightSessionGroups.get(from.sessionId);
+		if (groupId === undefined) {
+			return;
+		}
+		this._inFlightSessionGroups.delete(from.sessionId);
+		this._inFlightSessionGroups.set(to.sessionId, groupId);
+		this._onDidChange.fire({ groupsChanged: false, membershipChanged: new Set([from.sessionId, to.sessionId]) });
+	}
+
+	private setCurrentMembership(sessionId: string, groupId: string, changed: Set<string>): void {
+		const inFlightGroupId = this._inFlightSessionGroups.get(sessionId);
+		if (inFlightGroupId !== undefined) {
+			if (inFlightGroupId !== groupId) {
+				this._inFlightSessionGroups.set(sessionId, groupId);
+				changed.add(sessionId);
+			}
+			return;
+		}
+		this.setMembership(sessionId, groupId, changed);
+	}
+
 	private setMembership(sessionId: string, groupId: string, changed: Set<string>): void {
 		if (this._explicitlyUngroupedSessionIds.delete(sessionId) || this._membership.get(sessionId) !== groupId) {
 			this._membership.set(sessionId, groupId);

src/vs/sessions/services/sessions/browser/sessionsManagementService.ts

@@ -923,9 +923,7 @@ export class SessionsManagementService extends Disposable implements ISessionsMa
 		if (requestActivity) {
 			requestActivity.value = preparedRequestActivity;
 		}
-		if (replaceCurrentDraft) {
-			this._onDidReplaceNewDraftSession.fire({ from: originalSession, to: preparedSession });
-		}
+		this._onDidReplaceNewDraftSession.fire({ from: originalSession, to: preparedSession });
 		return { provider: preparedProvider, session: preparedSession };
 	}
 

src/vs/sessions/services/sessions/common/sessionsManagement.ts

@@ -385,8 +385,8 @@ export interface ISessionsManagementService {
 	 */
 	readonly onDidDiscardNewSession: Event<ISession>;
 	/**
-	 * Fires when {@link createNewSession} replaces the current in-progress
-	 * draft with another new-session draft (New Session → New Session).
+	 * Fires when an in-progress draft is replaced before sending,
+	 * either by {@link createNewSession} or new-session preparation.
 	 * Draft graduation uses {@link onDidReplaceSession} instead.
 	 */
 	readonly onDidReplaceNewDraftSession: Event<{ readonly from: ISession; readonly to: ISession }>;

src/vs/sessions/services/sessions/test/browser/sessionGroupsService.test.ts

@@ -54,9 +54,11 @@ suite('SessionGroupsService', () => {
 	let sessionUnarchivedEmitter: Emitter<ISession>;
 	let sessionDeletedEmitter: Emitter<ISession>;
 	let sessionReplacedEmitter: Emitter<{ readonly from: ISession; readonly to: ISession }>;
+	let newDraftSessionReplacedEmitter: Emitter<{ readonly from: ISession; readonly to: ISession }>;
 	let newSessionDiscardedEmitter: Emitter<ISession>;
 	let instantiationService: TestInstantiationService;
 	let sessions: ISession[];
+	let inFlightSessions: ISession[];
 
 	/** Simulate a new-session send: dispatch (`onWillSendRequest`) then start. */
 	function sendNewSession(draftId: string, committedId: string = draftId): void {
@@ -78,11 +80,14 @@ suite('SessionGroupsService', () => {
 		sessionUnarchivedEmitter = disposables.add(new Emitter<ISession>());
 		sessionDeletedEmitter = disposables.add(new Emitter<ISession>());
 		sessionReplacedEmitter = disposables.add(new Emitter<{ readonly from: ISession; readonly to: ISession }>());
+		newDraftSessionReplacedEmitter = disposables.add(new Emitter<{ readonly from: ISession; readonly to: ISession }>());
 		newSessionDiscardedEmitter = disposables.add(new Emitter<ISession>());
 		sessions = [];
+		inFlightSessions = [];
 		instantiationService.stub(ISessionsManagementService, {
 			...mock<ISessionsManagementService>(),
 			getSessions: () => sessions,
+			getInFlightNewSessionRequests: () => inFlightSessions,
 			getSession: resource => sessions.find(session => session.resource.toString() === resource.toString()),
 			onDidChangeSessions: sessionsChangedEmitter.event,
 			onWillSendRequest: willSendRequestEmitter.event,
@@ -91,6 +96,7 @@ suite('SessionGroupsService', () => {
 			onDidUnarchiveSession: sessionUnarchivedEmitter.event,
 			onDidDeleteSession: sessionDeletedEmitter.event,
 			onDidReplaceSession: sessionReplacedEmitter.event,
+			onDidReplaceNewDraftSession: newDraftSessionReplacedEmitter.event,
 			onDidDiscardNewSession: newSessionDiscardedEmitter.event,
 		});
 		service = disposables.add(instantiationService.createInstance(SessionGroupsService));
@@ -598,6 +604,90 @@ suite('SessionGroupsService', () => {
 		assert.deepStrictEqual(service.getSessionIdsInGroup(a.id), ['started']);
 	});
 
+	test('pending group applies while a provisional new session is published', () => {
+		const a = service.createGroup('A');
+		const draft = createSession('draft');
+		const committed = createSession('committed');
+		service.setPendingNewSessionGroup(a.id);
+		sessions = [draft];
+		inFlightSessions = [draft];
+
+		sessionsChangedEmitter.fire({ added: [draft], removed: [], changed: [] });
+
+		assert.strictEqual(service.getGroupOfSession(draft.sessionId), a.id);
+		assert.deepStrictEqual(service.getSessionIdsInGroup(a.id), [draft.sessionId]);
+
+		newDraftSessionReplacedEmitter.fire({ from: draft, to: committed });
+		sessions = [committed];
+		sessionsChangedEmitter.fire({ added: [], removed: [draft], changed: [committed] });
+		willSendRequestEmitter.fire(committed);
+		sessionStartedEmitter.fire(committed);
+
+		assert.deepStrictEqual({
+			draftGroup: service.getGroupOfSession(draft.sessionId),
+			committedGroup: service.getGroupOfSession(committed.sessionId),
+			groupMembers: service.getSessionIdsInGroup(a.id),
+		}, {
+			draftGroup: undefined,
+			committedGroup: a.id,
+			groupMembers: [committed.sessionId],
+		});
+	});
+
+	test('failed provisional new session clears transient group placement', () => {
+		const a = service.createGroup('A');
+		const draft = createSession('draft');
+		service.setPendingNewSessionGroup(a.id);
+		sessions = [draft];
+		inFlightSessions = [draft];
+		sessionsChangedEmitter.fire({ added: [draft], removed: [], changed: [] });
+
+		sessions = [];
+		inFlightSessions = [];
+		sessionsChangedEmitter.fire({ added: [], removed: [], changed: [] });
+
+		assert.strictEqual(service.getGroupOfSession(draft.sessionId), undefined);
+		assert.deepStrictEqual(service.getSessionIdsInGroup(a.id), []);
+	});
+
+	test('moving a provisional session overrides its pending group', () => {
+		const a = service.createGroup('A');
+		const b = service.createGroup('B');
+		const draft = createSession('draft');
+		service.setPendingNewSessionGroup(a.id);
+		sessions = [draft];
+		inFlightSessions = [draft];
+		sessionsChangedEmitter.fire({ added: [draft], removed: [], changed: [] });
+
+		service.addToGroup(draft.sessionId, b.id);
+		sessionStartedEmitter.fire(draft);
+
+		assert.deepStrictEqual({
+			group: service.getGroupOfSession(draft.sessionId),
+			aMembers: service.getSessionIdsInGroup(a.id),
+			bMembers: service.getSessionIdsInGroup(b.id),
+		}, {
+			group: b.id,
+			aMembers: [],
+			bMembers: [draft.sessionId],
+		});
+	});
+
+	test('removing a provisional session prevents its pending group from returning', () => {
+		const a = service.createGroup('A');
+		const draft = createSession('draft');
+		service.setPendingNewSessionGroup(a.id);
+		sessions = [draft];
+		inFlightSessions = [draft];
+		sessionsChangedEmitter.fire({ added: [draft], removed: [], changed: [] });
+
+		service.removeFromGroup(draft.sessionId);
+		sessionStartedEmitter.fire(draft);
+
+		assert.strictEqual(service.getGroupOfSession(draft.sessionId), undefined);
+		assert.deepStrictEqual(service.getSessionIdsInGroup(a.id), []);
+	});
+
 	test('pending group follows the draft as it graduates to a committed id', () => {
 		const a = service.createGroup('A');
 		service.setPendingNewSessionGroup(a.id);

src/vs/sessions/services/sessions/test/browser/sessionsManagementService.test.ts

@@ -3246,7 +3246,7 @@ suite('SessionsManagementService', () => {
 			sentSessionId,
 		}, {
 			deleted: ['s1'],
-			replacements: [],
+			replacements: ['s1->s2'],
 			sentSessionId: 's2',
 		});
 		completeSendRequest?.();