Merge pull request #294754 from mjbvz/dev/mjbvz/minimal-sawfish

Fix restoring non-local chat session

src/vs/workbench/contrib/chat/common/chatService/chatServiceImpl.ts

@@ -437,15 +437,14 @@ export class ChatService extends Disposable implements IChatService {
 			initialData: undefined,
 			location,
 			sessionResource,
-			sessionId,
 			canUseTools: options?.canUseTools ?? true,
 			disableBackgroundKeepAlive: options?.disableBackgroundKeepAlive
 		});
 	}
 
 	private _startSession(props: IStartSessionProps): ChatModel {
-		const { initialData, location, sessionResource, sessionId, canUseTools, transferEditingSession, disableBackgroundKeepAlive, inputState } = props;
-		const model = this.instantiationService.createInstance(ChatModel, initialData, { initialLocation: location, canUseTools, resource: sessionResource, sessionId, disableBackgroundKeepAlive, inputState });
+		const { initialData, location, sessionResource, canUseTools, transferEditingSession, disableBackgroundKeepAlive, inputState } = props;
+		const model = this.instantiationService.createInstance(ChatModel, initialData, { initialLocation: location, canUseTools, resource: sessionResource, disableBackgroundKeepAlive, inputState });
 		if (location === ChatAgentLocation.Chat) {
 			model.startEditingSession(true, transferEditingSession);
 		}
@@ -503,17 +502,17 @@ export class ChatService extends Disposable implements IChatService {
 			return existingRef;
 		}
 
-		const sessionId = LocalChatSessionUri.parseLocalSessionId(sessionResource);
-		if (!sessionId) {
-			throw new Error(`Cannot restore non-local session ${sessionResource}`);
-		}
-
 		let sessionData: ISerializedChatDataReference | undefined;
 		if (isEqual(this.transferredSessionResource, sessionResource)) {
 			this._transferredSessionResource = undefined;
 			sessionData = await this._chatSessionStore.readTransferredSession(sessionResource);
 		} else {
-			sessionData = await this._chatSessionStore.readSession(sessionId);
+			const localSessionId = LocalChatSessionUri.parseLocalSessionId(sessionResource);
+			if (localSessionId) {
+				sessionData = await this._chatSessionStore.readSession(localSessionId);
+			} else {
+				return this.loadSessionForResource(sessionResource, ChatAgentLocation.Chat, CancellationToken.None);
+			}
 		}
 
 		if (!sessionData) {
@@ -524,7 +523,6 @@ export class ChatService extends Disposable implements IChatService {
 			initialData: sessionData,
 			location: sessionData.value.initialLocation ?? ChatAgentLocation.Chat,
 			sessionResource,
-			sessionId,
 			canUseTools: true,
 		});
 
@@ -550,7 +548,6 @@ export class ChatService extends Disposable implements IChatService {
 			initialData: { value: data, serializer: new ChatSessionOperationLog() },
 			location: data.initialLocation ?? ChatAgentLocation.Chat,
 			sessionResource,
-			sessionId,
 			canUseTools: true,
 		});
 	}
@@ -568,6 +565,14 @@ export class ChatService extends Disposable implements IChatService {
 		}
 
 		const providedSession = await this.chatSessionService.getOrCreateChatSession(chatSessionResource, CancellationToken.None);
+
+		// Make sure we haven't created this in the meantime
+		const existingRefAfterProvision = this._sessionModels.acquireExisting(chatSessionResource);
+		if (existingRefAfterProvision) {
+			providedSession.dispose();
+			return existingRefAfterProvision;
+		}
+
 		const chatSessionType = chatSessionResource.scheme;
 
 		// Contributed sessions do not use UI tools

src/vs/workbench/contrib/chat/common/model/chatModel.ts

@@ -37,7 +37,7 @@ import { IChatEditingService, IChatEditingSession, ModifiedFileEntryState } from
 import { ILanguageModelChatMetadata, ILanguageModelChatMetadataAndIdentifier } from '../languageModels.js';
 import { IChatAgentCommand, IChatAgentData, IChatAgentResult, IChatAgentService, UserSelectedTools, reviveSerializedAgent } from '../participants/chatAgents.js';
 import { ChatRequestTextPart, IParsedChatRequest, reviveParsedChatRequest } from '../requestParser/chatParserTypes.js';
-import { LocalChatSessionUri } from './chatUri.js';
+import { chatSessionResourceToId, LocalChatSessionUri } from './chatUri.js';
 import { ObjectMutationLog } from './objectMutationLog.js';
 
 
@@ -2102,7 +2102,7 @@ export class ChatModel extends Disposable implements IChatModel {
 
 	constructor(
 		dataRef: ISerializedChatDataReference | undefined,
-		initialModelProps: { initialLocation: ChatAgentLocation; canUseTools: boolean; inputState?: ISerializableChatModelInputState; resource?: URI; sessionId?: string; disableBackgroundKeepAlive?: boolean },
+		initialModelProps: { initialLocation: ChatAgentLocation; canUseTools: boolean; inputState?: ISerializableChatModelInputState; resource?: URI; disableBackgroundKeepAlive?: boolean },
 		@ILogService private readonly logService: ILogService,
 		@IChatAgentService private readonly chatAgentService: IChatAgentService,
 		@IChatEditingService private readonly chatEditingService: IChatEditingService,
@@ -2118,8 +2118,23 @@ export class ChatModel extends Disposable implements IChatModel {
 		}
 
 		this._isImported = !!initialData && isValidExportedData && !isValidFullData;
-		this._sessionId = (isValidFullData && initialData.sessionId) || initialModelProps.sessionId || generateUuid();
-		this._sessionResource = initialModelProps.resource ?? LocalChatSessionUri.forSession(this._sessionId);
+
+		// Set the session resource and id
+		if (initialModelProps.resource) {
+			// prefer using the provided resource if provided
+			this._sessionId = chatSessionResourceToId(initialModelProps.resource);
+			this._sessionResource = initialModelProps.resource;
+		} else if (isValidFullData) {
+			// Otherwise use the serialized id. This is only valid for local chat sessions
+			this._sessionId = initialData.sessionId;
+			this._sessionResource = LocalChatSessionUri.forSession(initialData.sessionId);
+		} else {
+			// Finally fall back to generating a new id for a local session. This is used in the case where a
+			// chat has been exported (but not serialized)
+			this._sessionId = generateUuid();
+			this._sessionResource = LocalChatSessionUri.forSession(this._sessionId);
+		}
+
 		this._disableBackgroundKeepAlive = initialModelProps.disableBackgroundKeepAlive ?? false;
 
 		this._requests = initialData ? this._deserialize(initialData) : [];

src/vs/workbench/contrib/chat/common/model/chatModelStore.ts

@@ -16,7 +16,6 @@ export interface IStartSessionProps {
 	readonly initialData?: ISerializedChatDataReference;
 	readonly location: ChatAgentLocation;
 	readonly sessionResource: URI;
-	readonly sessionId?: string;
 	readonly canUseTools: boolean;
 	readonly transferEditingSession?: IChatEditingSession;
 	readonly disableBackgroundKeepAlive?: boolean;