From f4602885bdbfd67dccfabdb25cb0519b45912f56 Mon Sep 17 00:00:00 2001 From: Matt Bierner <12821956+mjbvz@users.noreply.github.com> Date: Thu, 12 Mar 2026 17:27:04 -0700 Subject: [PATCH] Move untitled session mapping much earlier in request process Seeing if we can avoid adding more workaround by making sure the swap happens much earlier in the process --- .../api/browser/mainThreadChatAgents2.ts | 30 ++----------- .../chatSessions/chatSessions.contribution.ts | 15 +++++++ .../contrib/chat/browser/widget/chatWidget.ts | 9 ++++ .../chat/common/chatService/chatService.ts | 2 + .../common/chatService/chatServiceImpl.ts | 45 +++++++++++++++++-- .../chat/common/chatSessionsService.ts | 1 + .../test/common/mockChatSessionsService.ts | 5 +++ 7 files changed, 76 insertions(+), 31 deletions(-) diff --git a/src/vs/workbench/api/browser/mainThreadChatAgents2.ts b/src/vs/workbench/api/browser/mainThreadChatAgents2.ts index 5924c301f0d..f40d952f332 100644 --- a/src/vs/workbench/api/browser/mainThreadChatAgents2.ts +++ b/src/vs/workbench/api/browser/mainThreadChatAgents2.ts @@ -42,7 +42,7 @@ import { IExtensionService } from '../../services/extensions/common/extensions.j import { Dto } from '../../services/extensions/common/proxyIdentifier.js'; import { ExtHostChatAgentsShape2, ExtHostContext, IChatNotebookEditDto, IChatParticipantMetadata, IChatProgressDto, IChatSessionContextDto, ICustomAgentDto, IDynamicChatAgentProps, IExtensionChatAgentMetadata, IInstructionDto, ISkillDto, MainContext, MainThreadChatAgentsShape2 } from '../common/extHost.protocol.js'; import { NotebookDto } from './mainThreadNotebookDto.js'; -import { getChatSessionType, isUntitledChatSession } from '../../contrib/chat/common/model/chatUri.js'; +import { isUntitledChatSession } from '../../contrib/chat/common/model/chatUri.js'; interface AgentData { dispose: () => void; @@ -247,32 +247,8 @@ export class MainThreadChatAgents2 extends Disposable implements MainThreadChatA const contributedSession = chatSession?.contributedChatSession; let chatSessionContext: IChatSessionContextDto | undefined; if (contributedSession) { - let chatSessionResource = contributedSession.chatSessionResource; - let isUntitled = isUntitledChatSession(chatSessionResource); - - // For new untitled sessions, invoke the controller's newChatSessionItemHandler - // to let the extension create a proper session item before the first request. - if (isUntitled) { - const newItem = await this._chatSessionService.createNewChatSessionItem(getChatSessionType(contributedSession.chatSessionResource), { prompt: request.message, command: request.command }, token); - if (newItem) { - chatSessionResource = newItem.resource; - isUntitled = false; - - // Update the model's contributed session with the resolved resource - // so subsequent requests don't re-invoke newChatSessionItemHandler - // and getChatSessionFromInternalUri returns the real resource. - chatSession?.setContributedChatSession({ - chatSessionResource, - initialSessionOptions: contributedSession.initialSessionOptions, - }); - - // Register alias so session-option lookups work with the new resource - this._chatSessionService.registerSessionResourceAlias( - contributedSession.chatSessionResource, - chatSessionResource - ); - } - } + const chatSessionResource = contributedSession.chatSessionResource; + const isUntitled = isUntitledChatSession(chatSessionResource); chatSessionContext = { chatSessionResource, diff --git a/src/vs/workbench/contrib/chat/browser/chatSessions/chatSessions.contribution.ts b/src/vs/workbench/contrib/chat/browser/chatSessions/chatSessions.contribution.ts index 5e80dea19d5..ee72554e7de 100644 --- a/src/vs/workbench/contrib/chat/browser/chatSessions/chatSessions.contribution.ts +++ b/src/vs/workbench/contrib/chat/browser/chatSessions/chatSessions.contribution.ts @@ -234,6 +234,9 @@ class ContributedChatSessionData extends Disposable { public getOption(optionId: string): string | IChatSessionProviderOptionItem | undefined { return this._optionsCache.get(optionId); } + public getAllOptions(): IterableIterator<[string, string | IChatSessionProviderOptionItem]> { + return this._optionsCache.entries(); + } public setOption(optionId: string, value: string | IChatSessionProviderOptionItem): void { this._optionsCache.set(optionId, value); } @@ -1073,6 +1076,18 @@ export class ChatSessionsService extends Disposable implements IChatSessionsServ return !!session && !!session.options && Object.keys(session.options).length > 0; } + public getSessionOptions(sessionResource: URI): Map | undefined { + const session = this._sessions.get(this._resolveResource(sessionResource)); + if (!session) { + return undefined; + } + const result = new Map(); + for (const [key, value] of session.getAllOptions()) { + result.set(key, typeof value === 'string' ? value : value.id); + } + return result.size > 0 ? result : undefined; + } + public getSessionOption(sessionResource: URI, optionId: string): string | IChatSessionProviderOptionItem | undefined { const session = this._sessions.get(this._resolveResource(sessionResource)); return session?.getOption(optionId); diff --git a/src/vs/workbench/contrib/chat/browser/widget/chatWidget.ts b/src/vs/workbench/contrib/chat/browser/widget/chatWidget.ts index 056229ddb23..d79fcf705f9 100644 --- a/src/vs/workbench/contrib/chat/browser/widget/chatWidget.ts +++ b/src/vs/workbench/contrib/chat/browser/widget/chatWidget.ts @@ -2378,6 +2378,15 @@ export class ChatWidget extends Disposable implements IChatWidget { this._onDidSubmitAgent.fire({ agent: sent.data.agent, slashCommand: sent.data.slashCommand }); this.handleDelegationExitIfNeeded(this._lockedAgent, sent.data.agent); + + // If the session was replaced (untitled -> real contributed session), swap the widget's model + if (sent.newSessionResource) { + const newModel = this.chatService.getSession(sent.newSessionResource); + if (newModel) { + this.setModel(newModel); + } + } + sent.data.responseCreatedPromise.then(() => { // Only start accessibility progress once a real request/response model exists. this.chatAccessibilityService.acceptRequest(submittedSessionResource); diff --git a/src/vs/workbench/contrib/chat/common/chatService/chatService.ts b/src/vs/workbench/contrib/chat/common/chatService/chatService.ts index bd161cf9098..102f0aefd5a 100644 --- a/src/vs/workbench/contrib/chat/common/chatService/chatService.ts +++ b/src/vs/workbench/contrib/chat/common/chatService/chatService.ts @@ -1288,6 +1288,8 @@ export interface ChatSendResultRejected { export interface ChatSendResultSent { readonly kind: 'sent'; readonly data: IChatSendRequestData; + /** Set when the session was replaced by a new one (e.g. untitled -> real contributed session). */ + readonly newSessionResource?: URI; } export interface ChatSendResultQueued { diff --git a/src/vs/workbench/contrib/chat/common/chatService/chatServiceImpl.ts b/src/vs/workbench/contrib/chat/common/chatService/chatServiceImpl.ts index d00dff8c817..32fee250be1 100644 --- a/src/vs/workbench/contrib/chat/common/chatService/chatServiceImpl.ts +++ b/src/vs/workbench/contrib/chat/common/chatService/chatServiceImpl.ts @@ -42,11 +42,11 @@ import { chatAgentLeader, ChatRequestAgentPart, ChatRequestAgentSubcommandPart, import { ChatRequestParser } from '../requestParser/chatRequestParser.js'; import { ChatMcpServersStarting, ChatPendingRequestChangeClassification, ChatPendingRequestChangeEvent, ChatPendingRequestChangeEventName, ChatRequestQueueKind, ChatSendResult, ChatSendResultQueued, ChatSendResultSent, ChatStopCancellationNoopClassification, ChatStopCancellationNoopEvent, ChatStopCancellationNoopEventName, IChatCompleteResponse, IChatDetail, IChatFollowup, IChatModelReference, IChatProgress, IChatQuestionAnswers, IChatSendRequestOptions, IChatSendRequestResponseState, IChatService, IChatSessionContext, IChatSessionStartOptions, IChatUserActionEvent, ResponseModelState } from './chatService.js'; import { ChatRequestTelemetry, ChatServiceTelemetry } from './chatServiceTelemetry.js'; -import { IChatSessionsService } from '../chatSessionsService.js'; +import { IChatSessionsService, localChatSessionType } from '../chatSessionsService.js'; import { ChatSessionStore, IChatSessionEntryMetadata } from '../model/chatSessionStore.js'; import { IChatSlashCommandService } from '../participants/chatSlashCommands.js'; import { IChatTransferService } from '../model/chatTransferService.js'; -import { chatSessionResourceToId, LocalChatSessionUri } from '../model/chatUri.js'; +import { chatSessionResourceToId, getChatSessionType, isUntitledChatSession, LocalChatSessionUri } from '../model/chatUri.js'; import { IChatRequestVariableEntry } from '../attachments/chatVariableEntries.js'; import { ChatAgentLocation, ChatModeKind } from '../constants.js'; import { ChatMessageRole, IChatMessage, ILanguageModelsService } from '../languageModels.js'; @@ -602,7 +602,7 @@ export class ChatService extends Disposable implements IChatService { } } - const chatSessionType = sessionResource.scheme; + const chatSessionType = getChatSessionType(sessionResource); // Contributed sessions do not use UI tools const modelRef = this._sessionModels.acquireOrCreate({ @@ -805,11 +805,47 @@ export class ChatService extends Disposable implements IChatService { return { kind: 'rejected', reason: 'Empty message' }; } - const model = this._sessionModels.get(sessionResource); + let model = this._sessionModels.get(sessionResource); if (!model) { throw new Error(`Unknown session: ${sessionResource}`); } + let newSessionResource: URI | undefined; + + // Workaround for the contributed chat sessions + // + // Internally blank widgets uses special sessions with an untitled- path. We do not want these leaking out + // to the rest of code. Instead use `createNewChatSessionItem` to make sure the session gets properly initialized with a real resource before processing the first request. + if (!model.hasRequests && isUntitledChatSession(sessionResource) && getChatSessionType(sessionResource) !== localChatSessionType) { + + const parsedRequest = this.parseChatRequest(sessionResource, request, options?.location ?? model.initialLocation, options); + const commandPart = parsedRequest.parts.find((r): r is ChatRequestSlashCommandPart => r instanceof ChatRequestSlashCommandPart); + const requestText = getPromptText(parsedRequest).message; + const newItem = await this.chatSessionService.createNewChatSessionItem(getChatSessionType(sessionResource), { prompt: requestText, command: commandPart?.text }, CancellationToken.None); + if (newItem) { + + // Update the model's contributed session with the resolved resource + // so subsequent requests don't re-invoke newChatSessionItemHandler + // and getChatSessionFromInternalUri returns the real resource. + const sessionOptions = this.chatSessionService.getSessionOptions(sessionResource); + model?.setContributedChatSession({ + chatSessionResource: sessionResource, + initialSessionOptions: sessionOptions ? [...sessionOptions].map(([optionId, value]) => ({ optionId, value })) : undefined, + }); + + model = (await this.loadRemoteSession(newItem.resource, model.initialLocation, CancellationToken.None))?.object as ChatModel | undefined; + if (!model) { + throw new Error(`Failed to load session for resource: ${newItem.resource}`); + } + + // Register alias so session-option lookups work with the new resource + this.chatSessionService.registerSessionResourceAlias(sessionResource, newItem.resource); + + sessionResource = newItem.resource; + newSessionResource = newItem.resource; + } + } + const hasPendingRequest = this._pendingRequests.has(sessionResource); if (options?.queue) { @@ -847,6 +883,7 @@ export class ChatService extends Disposable implements IChatService { // This method is only returning whether the request was accepted - don't block on the actual request return { kind: 'sent', + newSessionResource, data: { ...this._sendRequestAsync(model, sessionResource, parsedRequest, attempt, !options?.noCommandDetection, silentAgent ?? defaultAgent, location, options), agent, diff --git a/src/vs/workbench/contrib/chat/common/chatSessionsService.ts b/src/vs/workbench/contrib/chat/common/chatSessionsService.ts index 01020f2b624..1b843bce36b 100644 --- a/src/vs/workbench/contrib/chat/common/chatSessionsService.ts +++ b/src/vs/workbench/contrib/chat/common/chatSessionsService.ts @@ -274,6 +274,7 @@ export interface IChatSessionsService { getOrCreateChatSession(sessionResource: URI, token: CancellationToken): Promise; hasAnySessionOptions(sessionResource: URI): boolean; + getSessionOptions(sessionResource: URI): Map | undefined; getSessionOption(sessionResource: URI, optionId: string): string | IChatSessionProviderOptionItem | undefined; setSessionOption(sessionResource: URI, optionId: string, value: string | IChatSessionProviderOptionItem): boolean; diff --git a/src/vs/workbench/contrib/chat/test/common/mockChatSessionsService.ts b/src/vs/workbench/contrib/chat/test/common/mockChatSessionsService.ts index 754091ae0c1..bae11e420fb 100644 --- a/src/vs/workbench/contrib/chat/test/common/mockChatSessionsService.ts +++ b/src/vs/workbench/contrib/chat/test/common/mockChatSessionsService.ts @@ -181,6 +181,11 @@ export class MockChatSessionsService implements IChatSessionsService { await this._onRequestNotifyExtension.fireAsync({ sessionResource, updates }, CancellationToken.None); } + getSessionOptions(sessionResource: URI): Map | undefined { + const options = this.sessionOptions.get(sessionResource); + return options && options.size > 0 ? options : undefined; + } + getSessionOption(sessionResource: URI, optionId: string): string | undefined { return this.sessionOptions.get(sessionResource)?.get(optionId); }