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
This commit is contained in:
Matt Bierner
2026-03-12 17:27:04 -07:00
parent d9195789ea
commit f4602885bd
7 changed files with 76 additions and 31 deletions

View File

@@ -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,

View File

@@ -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<string, string> | undefined {
const session = this._sessions.get(this._resolveResource(sessionResource));
if (!session) {
return undefined;
}
const result = new Map<string, string>();
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);

View File

@@ -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);

View File

@@ -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 {

View File

@@ -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,

View File

@@ -274,6 +274,7 @@ export interface IChatSessionsService {
getOrCreateChatSession(sessionResource: URI, token: CancellationToken): Promise<IChatSession>;
hasAnySessionOptions(sessionResource: URI): boolean;
getSessionOptions(sessionResource: URI): Map<string, string> | undefined;
getSessionOption(sessionResource: URI, optionId: string): string | IChatSessionProviderOptionItem | undefined;
setSessionOption(sessionResource: URI, optionId: string, value: string | IChatSessionProviderOptionItem): boolean;

View File

@@ -181,6 +181,11 @@ export class MockChatSessionsService implements IChatSessionsService {
await this._onRequestNotifyExtension.fireAsync({ sessionResource, updates }, CancellationToken.None);
}
getSessionOptions(sessionResource: URI): Map<string, string> | 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);
}