Skip to content

[cherry-pick] agentHost: Refine workspace transition boundaries #334947

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Sandeep Somavarapu (sandy081) merged 1 commit into from
Sep 7, 2026
Merged

[cherry-pick] agentHost: Refine workspace transition boundaries #334947

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/vs/platform/agentHost/common/agentHostChatContributionsService.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,8 +108,8 @@ export type IncomingRequestDisposition =
export interface IHydrationContext {
readonly session: ProtocolURI;
readonly chat: ProtocolURI;
/** Authoritative value from already-loaded host session metadata, when available. */
readonly hasWorkspaceTransitions?: boolean;
/** Workspace transitions preloaded from this chat's backing storage, when present. */
readonly workspaceTransitions?: ReadonlyMap<string, string>;
}

/**
Expand Down
60 changes: 48 additions & 12 deletions src/vs/platform/agentHost/node/agentService.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ import { decodeBase64, encodeBase64, VSBuffer } from '../../../base/common/buffe
import { Barrier, DeferredPromise, disposableTimeout, Limiter, ResourceQueue } from '../../../base/common/async.js';
import { toErrorMessage } from '../../../base/common/errorMessage.js';
import { Emitter } from '../../../base/common/event.js';
import { Disposable, DisposableMap, DisposableResourceMap, DisposableStore, IDisposable, MutableDisposable, toDisposable } from '../../../base/common/lifecycle.js';
import { Disposable, DisposableMap, DisposableResourceMap, DisposableStore, IDisposable, IReference, MutableDisposable, toDisposable } from '../../../base/common/lifecycle.js';
import { getExtensionForMimeType, getMediaMime, getMediaOrTextMime } from '../../../base/common/mime.js';
import { Schemas } from '../../../base/common/network.js';
import { dirname as resourcesDirname, extname as resourcesExtname, extUriBiasedIgnorePathCase, isEqual, isEqualOrParent, joinPath } from '../../../base/common/resources.js';
Expand All @@ -21,7 +21,7 @@ import { IInstantiationService } from '../../instantiation/common/instantiation.
import { ILogService } from '../../log/common/log.js';
import { AgentChatMigrationDeferred, AgentProvider, AgentSession, AgentSignal, IAgent, type IAgentAdoptedWorktree, IAgentChatContext, IAgentChatDataChange, IAgentChatMetadata, IAgentCreateChatOptions, IAgentCreateChatRequestOptions, IAgentCreateChatResult, IAgentCreateChatSideChatSelection, IAgentCreateChatSideChatSource, IAgentCreateSessionConfig, IAgentCreateSessionResult, IAgentDiscoveredChat, IAgentMaterializeChatEvent, IAgentModelInfo, IAgentResolveSessionConfigParams, IAgentChatAdoptionResult, type AgentChatAdoptionReason, IAgentSessionConfigCompletionsParams, IAgentSessionMetadata, IAgentSpawnChatEvent, AuthenticateParams, AuthenticateResult, SubagentChatSignal, subagentChatTitle } from '../common/agent.js';
import { type AgentHostDebugLogsArtifactKind, type IAgentHostDebugLogsArtifact, type IAgentHostDebugLogsChunk, IAgentHostManagedSettingsDiagnostics, IAgentHostNetworkDiagnosticsInfo, IAgentHostNetworkFetchResult, IAgentService } from '../common/agentService.js';
import { ISessionDataService, SESSION_ATTACHMENTS_DIRNAME } from '../common/sessionDataService.js';
import { ISessionDatabase, ISessionDataService, SESSION_ATTACHMENTS_DIRNAME } from '../common/sessionDataService.js';
import { IAgentEditAttributionService, ICancelEditAttributionFlushParams, ICommitEditAttributionFlushParams, IEditAttributionFlushResult, IPrepareEditAttributionFlushParams, IPreparedEditAttributionFlush, parseEditAttributionResource } from '../common/fileEditAttribution.js';
import { omitTransientSessionConfigValues, SessionConfigKey } from '../common/sessionConfigKeys.js';
import type { IAgentCustomizationSettingsRegistration } from '../common/agentCustomizationSettings.js';
Expand All @@ -37,7 +37,7 @@ import type { InvokeChangesetOperationParams, InvokeChangesetOperationResult } f
import { AhpErrorCodes, AHP_SESSION_NOT_FOUND, ContentEncoding, JSON_RPC_INTERNAL_ERROR, ProtocolError, ResourceChangeType, ResourceType, ResourceWriteMode, type CreateResourceWatchParams, type CreateResourceWatchResult, type DirectoryEntry, type ResourceCopyParams, type ResourceCopyResult, type ResourceDeleteParams, type ResourceDeleteResult, type ResourceListResult, type ResourceMkdirParams, type ResourceMkdirResult, type ResourceMoveParams, type ResourceMoveResult, type ResourceReadResult, type ResourceResolveParams, type ResourceResolveResult, type ResourceWatchState, type ResourceWriteParams, type ResourceWriteResult, type IStateSnapshot } from '../common/state/sessionProtocol.js';
import { ChangesSummary, ChatInteractivity, ChatOriginKind, MessageAttachmentKind, type Annotation, type AnnotationEntry, type AnnotationOrigin, type AnnotationsState, type ChatOrigin, type Customization, type Message, type MessageAttachment, type MessageResourceAttachment, type TextRange } from '../common/state/protocol/state.js';
import type { ChatPendingMessageSetAction, ChatTurnStartedAction, SessionConfigChangedAction } from '../common/state/protocol/actions.js';
import { isAhpAutomationCatalogChannel, isAhpAutomationRunChannel, ISessionGitHubState, ISessionGitState, MessageKind, ResponsePartKind, SESSION_META_GITHUB_KEY, SESSION_META_GIT_KEY, SESSION_META_MULTI_ROOT_KEY, SESSION_META_SOURCE_CONTROL_KEY, AH_META_CREATED_BY_SESSION_DB_KEY, readSessionCreationReference, readSessionSpawnDepth, withSessionSpawnDepth, withSessionCreationReference, parseSessionCreationReference, SessionLifecycle, SessionStatus, ToolCallStatus, ToolResultContentType, TurnState, AH_META_HAS_WORKSPACE_TRANSITIONS_DB_KEY, AH_META_WORKSPACE_CONVERSION_QUARANTINED_DB_KEY, AH_META_WORKSPACELESS_DB_KEY, AH_META_EHCLI_ADOPTED_DB_KEY, AH_META_IS_ARCHIVED_DB_KEY, AH_META_IS_DONE_DB_KEY, AH_META_IS_READ_DB_KEY, buildChatUri, buildDefaultChatUri, buildResourceWatchChannelUri, buildSubagentChatUri, buildSubagentSessionUriPrefix, getErrorResponsePart, isAhpChatChannel, isChatReadOnly, isDefaultChatUri, isSubagentChatUri, isSubagentSession, needsSessionGitStateRefresh, parseChatUri, parseDefaultChatUri, parseRequiredSessionUriFromChatUri, parseResourceWatchChannelUri, parseSessionMultiRootMetadata, parseSubagentSessionUri, readSessionExternal, readSessionGitHubState, readSessionGitState, readSessionHasWorkspaceTransitions, readSessionMultiRootMetadata, readSessionSourceControlState, readSessionWorkspaceless, withMessageRequestHiddenFromTranscript, withSessionExternal, withSessionGitHubState, withSessionGitState, withSessionHasWorkspaceTransitions, withSessionMultiRootMetadata, withSessionSourceControlState, withSessionStatusFlag, withSessionWorkspaceless, withSessionEhcliAdopted, withSessionEhcliLastMigratedTurn, AH_META_EHCLI_LAST_TURN_DB_KEY, withSessionFolderPickerDecision, readSessionFolderPickerDecision, parseSessionFolderPickerDecision, SESSION_META_FOLDER_PICKER_KEY, readSessionEhcliAdoptable, type ISessionSourceControlState, type SessionConfigState, type SessionSummary, type ToolResultSubagentContent, type Turn } from '../common/state/sessionState.js';
import { isAhpAutomationCatalogChannel, isAhpAutomationRunChannel, ISessionGitHubState, ISessionGitState, MessageKind, ResponsePartKind, SESSION_META_GITHUB_KEY, SESSION_META_GIT_KEY, SESSION_META_MULTI_ROOT_KEY, SESSION_META_SOURCE_CONTROL_KEY, AH_META_CREATED_BY_SESSION_DB_KEY, readSessionCreationReference, readSessionSpawnDepth, withSessionSpawnDepth, withSessionCreationReference, parseSessionCreationReference, SessionLifecycle, SessionStatus, ToolCallStatus, ToolResultContentType, TurnState, AH_META_HAS_WORKSPACE_TRANSITIONS_DB_KEY, AH_META_WORKSPACE_CONVERSION_QUARANTINED_DB_KEY, AH_META_WORKSPACELESS_DB_KEY, AH_META_EHCLI_ADOPTED_DB_KEY, AH_META_IS_ARCHIVED_DB_KEY, AH_META_IS_DONE_DB_KEY, AH_META_IS_READ_DB_KEY, buildChatUri, buildDefaultChatUri, buildResourceWatchChannelUri, buildSubagentChatUri, buildSubagentSessionUriPrefix, chatStorageUri, getErrorResponsePart, isAhpChatChannel, isChatReadOnly, isDefaultChatUri, isSubagentChatUri, isSubagentSession, needsSessionGitStateRefresh, parseChatUri, parseDefaultChatUri, parseRequiredSessionUriFromChatUri, parseResourceWatchChannelUri, parseSessionMultiRootMetadata, parseSubagentSessionUri, readSessionExternal, readSessionGitHubState, readSessionGitState, readSessionMultiRootMetadata, readSessionSourceControlState, readSessionWorkspaceless, withMessageRequestHiddenFromTranscript, withSessionExternal, withSessionGitHubState, withSessionGitState, withSessionHasWorkspaceTransitions, withSessionMultiRootMetadata, withSessionSourceControlState, withSessionStatusFlag, withSessionWorkspaceless, withSessionEhcliAdopted, withSessionEhcliLastMigratedTurn, AH_META_EHCLI_LAST_TURN_DB_KEY, withSessionFolderPickerDecision, readSessionFolderPickerDecision, parseSessionFolderPickerDecision, SESSION_META_FOLDER_PICKER_KEY, readSessionEhcliAdoptable, type ISessionSourceControlState, type SessionConfigState, type SessionSummary, type ToolResultSubagentContent, type Turn } from '../common/state/sessionState.js';
import { readToolCallMeta } from '../common/meta/agentToolCallMeta.js';
import { isHostSnapshotAttachment, toHostSnapshotAttachmentMeta } from '../common/meta/agentSnapshotAttachmentMeta.js';
import { readEphemeralSessionMeta, withEphemeralSessionMeta } from '../common/meta/agentEphemeralSessionMeta.js';
Expand Down Expand Up @@ -3492,12 +3492,44 @@ export class AgentService extends Disposable implements IAgentService {
* exactly that window; every other caller relies on the exhaustive origin
* {@link _chatContext} stamps.
*/
private async _getChatMessages(provider: IAgent, chat: URI, session: URI, origin?: ChatOrigin, hasWorkspaceTransitions?: boolean): Promise<readonly Turn[]> {
private async _getChatMessages(provider: IAgent, chat: URI, session: URI, origin?: ChatOrigin, workspaceTransitionsPromise?: Promise<ReadonlyMap<string, string> | undefined>): Promise<readonly Turn[]> {
const context = { ...this._chatContext(session, chat), ...(origin ? { origin } : {}) };
const transitionsPromise = workspaceTransitionsPromise ?? this._loadWorkspaceTransitions(chat);
this._logService.trace(`[AgentService] getChatMessages start: chat=${chat.toString()}`);
const providerTurns = await provider.chats.getMessages(chat, context);
const [providerTurns, workspaceTransitions] = await Promise.all([
provider.chats.getMessages(chat, context),
transitionsPromise,
]);
this._logService.trace(`[AgentService] getChatMessages: provider returned ${providerTurns.length} turn(s) for chat=${chat.toString()}`);
return this._chatContributions.hydrateTurns({ session: session.toString(), chat: chat.toString(), hasWorkspaceTransitions }, providerTurns);
return this._chatContributions.hydrateTurns({ session: session.toString(), chat: chat.toString(), workspaceTransitions }, providerTurns);
}

private async _loadWorkspaceTransitions(chat: URI): Promise<ReadonlyMap<string, string> | undefined> {
const storage = chatStorageUri(chat);
if (!storage) {
return undefined;
}
try {
const database = await this._sessionDataService.tryOpenDatabase(storage);
return database ? this._readWorkspaceTransitions(database, storage, true) : undefined;
} catch (error) {
this._logService.warn(`[AgentService] Failed to open workspace transition storage for ${storage.toString()}`, error);
return undefined;
}
}

private async _readWorkspaceTransitions(database: IReference<ISessionDatabase>, storage: URI, requireMarker = false): Promise<ReadonlyMap<string, string> | undefined> {
try {
if (requireMarker && await database.object.getMetadata(AH_META_HAS_WORKSPACE_TRANSITIONS_DB_KEY) !== 'true') {
return undefined;
}
return await database.object.getTurnWorkspaceTransitions();
} catch (error) {
this._logService.warn(`[AgentService] Failed to restore workspace transitions for ${storage.toString()}`, error);
return undefined;
} finally {
database.dispose();
}
}

/**
Expand Down Expand Up @@ -5509,12 +5541,13 @@ export class AgentService extends Disposable implements IAgentService {
let gitMetadata: Record<string, string | undefined> | undefined;
let changesetMetadata: Record<string, string | undefined> | undefined;
let sessionMetadata: Record<string, unknown> | undefined;
let sessionMetadataRead = false;
let workspaceTransitionsPromise: Promise<ReadonlyMap<string, string> | undefined> = Promise.resolve(undefined);
const ref = this._sessionDataService.tryOpenDatabase?.(session);
if (ref) {
try {
const db = await ref;
if (db) {
let retainedForWorkspaceTransitions = false;
try {
const m = await db.object.getMetadataObject({
customTitle: true,
Expand All @@ -5534,7 +5567,10 @@ export class AgentService extends Disposable implements IAgentService {
...GIT_DB_METADATA_KEYS,
...CHANGESET_DB_METADATA_KEYS,
});
sessionMetadataRead = true;
if (m[AH_META_HAS_WORKSPACE_TRANSITIONS_DB_KEY] === 'true') {
retainedForWorkspaceTransitions = true;
workspaceTransitionsPromise = this._readWorkspaceTransitions(db, session);
}
if (m.customTitle) {
title = m.customTitle;
}
Expand Down Expand Up @@ -5626,10 +5662,10 @@ export class AgentService extends Disposable implements IAgentService {
}
}
} finally {
db.dispose();
if (!retainedForWorkspaceTransitions) {
db.dispose();
}
}
} else {
sessionMetadataRead = true;
}
} catch {
// Best-effort: fall back to agent-provided metadata
Expand All @@ -5639,7 +5675,7 @@ export class AgentService extends Disposable implements IAgentService {

let turns: readonly Turn[];
try {
turns = await this._getChatMessages(agent, defaultChatUri, session, undefined, sessionMetadataRead ? readSessionHasWorkspaceTransitions(sessionMetadata) : undefined);
turns = await this._getChatMessages(agent, defaultChatUri, session, undefined, workspaceTransitionsPromise);
} catch (err) {
if (err instanceof ProtocolError) {
throw err;
Expand Down
43 changes: 4 additions & 39 deletions ...de/chatContributions/sessionWorkspaceConversion/sessionWorkspaceConversionContribution.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,15 +4,11 @@
*--------------------------------------------------------------------------------------------*/

import { Disposable } from '../../../../../base/common/lifecycle.js';
import { URI } from '../../../../../base/common/uri.js';
import { localize } from '../../../../../nls.js';
import { ILogService } from '../../../../log/common/log.js';
import type { IAgentHostChatContribution, IAgentHostChatContributionContext, IHydrationContext, IIncomingRequest, IncomingRequestDisposition, ITurnEnd } from '../../../common/agentHostChatContributionsService.js';
import { parseAgentWorkspaceTransition, AgentSystemNotificationKind, readAgentSystemNotificationMeta, toAgentSystemNotificationMeta } from '../../../common/meta/agentSystemNotificationMeta.js';
import { toAgentWorkspaceContinuationMessageMeta } from '../../../common/meta/agentWorkspaceContinuationMeta.js';
import { ISessionDataService } from '../../../common/sessionDataService.js';
import { chatStorageUri, readSessionHasWorkspaceTransitions, ResponsePartKind, withMessageRequestHiddenFromTranscript, type Turn } from '../../../common/state/sessionState.js';
import { AgentHostStateManager, IAgentHostStateManager } from '../../agentHostStateManager.js';
import { ResponsePartKind, withMessageRequestHiddenFromTranscript, type Turn } from '../../../common/state/sessionState.js';
import { ISessionWorkspaceConversionService } from './sessionWorkspaceConversionService.js';

/** Finalizes requested workspace conversions after a turn and blocks new turns while conversion is pending. */
Expand All @@ -24,9 +20,6 @@ export class SessionWorkspaceConversionContribution extends Disposable implement
constructor(
protected readonly _context: IAgentHostChatContributionContext,
@ISessionWorkspaceConversionService private readonly _conversionService: ISessionWorkspaceConversionService,
@IAgentHostStateManager private readonly _stateManager: AgentHostStateManager,
@ISessionDataService private readonly _sessionDataService: ISessionDataService,
@ILogService private readonly _logService: ILogService,
) {
super();
}
Expand All @@ -53,37 +46,9 @@ export class SessionWorkspaceConversionContribution extends Disposable implement
};
}

async onHydrateTurns(context: IHydrationContext, turns: readonly Turn[]): Promise<readonly Turn[]> {
if (turns.length === 0) {
return turns;
}
if (context.hasWorkspaceTransitions === false) {
return turns;
}
if (context.hasWorkspaceTransitions !== true) {
const summary = this._stateManager.getSessionSummary(context.session);
if (summary && !readSessionHasWorkspaceTransitions(summary._meta)) {
return turns;
}
}
const storage = chatStorageUri(URI.parse(context.chat));
if (!storage) {
return turns;
}
const database = await this._sessionDataService.tryOpenDatabase(storage);
if (!database) {
return turns;
}
let transitions: Map<string, string>;
try {
transitions = await database.object.getTurnWorkspaceTransitions();
} catch (error) {
this._logService.warn(`[SessionWorkspaceConversionContribution] Failed to restore workspace transitions for ${storage.toString()}`, error);
return turns;
} finally {
database.dispose();
}
if (transitions.size === 0) {
onHydrateTurns(context: IHydrationContext, turns: readonly Turn[]): readonly Turn[] {
const transitions = context.workspaceTransitions;
if (turns.length === 0 || !transitions?.size) {
return turns;
}
return turns.map(turn => {
Expand Down
Loading