Skip to content

Commit 66b54e1

Browse files
jriekenjrieken
authored
joh/tasteless orca (microsoft#185006)
* Add ability to rerun prompt Extract `SessionPrompt` with tries counter that, add rerun message that is listened on during WAIT_FOR_INPUT, read number of attempts during MAKE_REQUEST microsoft/vscode-copilot#247 * renames
1 parent 3447686 commit 66b54e1

File tree

6 files changed

+115
-23
lines changed

6 files changed

+115
-23
lines changed

‎src/vs/workbench/contrib/inlineChat/browser/inlineChat.contribution.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ registerAction2(InlineChatActions.StartSessionAction);
2626
registerAction2(InlineChatActions.UnstashSessionAction);
2727
registerAction2(InlineChatActions.MakeRequestAction);
2828
registerAction2(InlineChatActions.StopRequestAction);
29+
registerAction2(InlineChatActions.ReRunRequestAction);
2930
registerAction2(InlineChatActions.DiscardAction);
3031
registerAction2(InlineChatActions.DiscardToClipboardAction);
3132
registerAction2(InlineChatActions.DiscardUndoToNewFileAction);

‎src/vs/workbench/contrib/inlineChat/browser/inlineChatActions.ts

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,28 @@ export class MakeRequestAction extends AbstractInlineChatAction {
159159
}
160160
}
161161

162+
export class ReRunRequestAction extends AbstractInlineChatAction {
163+
164+
constructor() {
165+
super({
166+
id: 'inlineChat.regenerate',
167+
title: localize('rerun', 'Regenerate Response'),
168+
icon: Codicon.refresh,
169+
precondition: ContextKeyExpr.and(CTX_INLINE_CHAT_VISIBLE, CTX_INLINE_CHAT_EMPTY.negate(), CTX_INLINE_CHAT_LAST_RESPONSE_TYPE),
170+
menu: {
171+
id: MENU_INLINE_CHAT_WIDGET_STATUS,
172+
group: '0_main',
173+
order: 1,
174+
}
175+
});
176+
}
177+
178+
override runInlineChatCommand(_accessor: ServicesAccessor, ctrl: InlineChatController): void {
179+
ctrl.regenerate();
180+
}
181+
182+
}
183+
162184
export class StopRequestAction extends AbstractInlineChatAction {
163185

164186
constructor() {

‎src/vs/workbench/contrib/inlineChat/browser/inlineChatController.ts

Lines changed: 24 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ import { IContextKey, IContextKeyService } from 'vs/platform/contextkey/common/c
2828
import { IDialogService } from 'vs/platform/dialogs/common/dialogs';
2929
import { IInstantiationService, ServicesAccessor } from 'vs/platform/instantiation/common/instantiation';
3030
import { ILogService } from 'vs/platform/log/common/log';
31-
import { EditResponse, EmptyResponse, ErrorResponse, ExpansionState, IInlineChatSessionService, MarkdownResponse, Session, SessionExchange } from 'vs/workbench/contrib/inlineChat/browser/inlineChatSession';
31+
import { EditResponse, EmptyResponse, ErrorResponse, ExpansionState, IInlineChatSessionService, MarkdownResponse, Session, SessionExchange, SessionPrompt } from 'vs/workbench/contrib/inlineChat/browser/inlineChatSession';
3232
import { EditModeStrategy, LivePreviewStrategy, LiveStrategy, PreviewStrategy } from 'vs/workbench/contrib/inlineChat/browser/inlineChatStrategies';
3333
import { InlineChatZoneWidget } from 'vs/workbench/contrib/inlineChat/browser/inlineChatWidget';
3434
import { CTX_INLINE_CHAT_HAS_ACTIVE_REQUEST, CTX_INLINE_CHAT_LAST_FEEDBACK, IInlineChatRequest, IInlineChatResponse, INLINE_CHAT_ID, EditMode, InlineChatResponseFeedbackKind, CTX_INLINE_CHAT_LAST_RESPONSE_TYPE, InlineChatResponseType, CTX_INLINE_CHAT_DID_EDIT, CTX_INLINE_CHAT_HAS_STASHED_SESSION } from 'vs/workbench/contrib/inlineChat/common/inlineChat';
@@ -58,7 +58,8 @@ const enum Message {
5858
PAUSE_SESSION = 1 << 2,
5959
CANCEL_REQUEST = 1 << 3,
6060
CANCEL_INPUT = 1 << 4,
61-
ACCEPT_INPUT = 1 << 5
61+
ACCEPT_INPUT = 1 << 5,
62+
RERUN_INPUT = 1 << 6,
6263
}
6364

6465
export interface InlineChatRunOptions {
@@ -267,7 +268,7 @@ export class InlineChatController implements IEditorContribution {
267268

268269
this._zone.value.widget.updateSlashCommands(this._activeSession.session.slashCommands ?? []);
269270
this._zone.value.widget.placeholder = this._getPlaceholderText();
270-
this._zone.value.widget.value = this._activeSession.lastInput ?? '';
271+
this._zone.value.widget.value = this._activeSession.lastInput?.value ?? '';
271272
this._zone.value.widget.updateInfo(this._activeSession.session.message ?? localize('welcome.1', "AI-generated code may be incorrect"));
272273
this._zone.value.show(this._activeSession.wholeRange.value.getEndPosition());
273274
this._zone.value.widget.preferredExpansionState = this._activeSession.lastExpansionState;
@@ -357,6 +358,7 @@ export class InlineChatController implements IEditorContribution {
357358

358359
private async [State.WAIT_FOR_INPUT](options: InlineChatRunOptions | undefined): Promise<State.ACCEPT | State.CANCEL | State.PAUSE | State.WAIT_FOR_INPUT | State.MAKE_REQUEST> {
359360
assertType(this._activeSession);
361+
assertType(this._strategy);
360362

361363
this._zone.value.widget.placeholder = this._getPlaceholderText();
362364
this._zone.value.show(this._activeSession.wholeRange.value.getEndPosition());
@@ -397,6 +399,15 @@ export class InlineChatController implements IEditorContribution {
397399
return State.PAUSE;
398400
}
399401

402+
if (message & Message.RERUN_INPUT && this._activeSession.lastExchange) {
403+
const { lastExchange } = this._activeSession;
404+
this._activeSession.addInput(lastExchange.prompt.retry());
405+
if (lastExchange.response instanceof EditResponse) {
406+
await this._strategy.undoChanges(lastExchange.response);
407+
}
408+
return State.MAKE_REQUEST;
409+
}
410+
400411
if (!this._zone.value.widget.value) {
401412
return State.WAIT_FOR_INPUT;
402413
}
@@ -420,7 +431,7 @@ export class InlineChatController implements IEditorContribution {
420431
return State.WAIT_FOR_INPUT;
421432
}
422433

423-
this._activeSession.addInput(input);
434+
this._activeSession.addInput(new SessionPrompt(input));
424435
return State.MAKE_REQUEST;
425436
}
426437

@@ -444,10 +455,10 @@ export class InlineChatController implements IEditorContribution {
444455

445456
const sw = StopWatch.create();
446457
const request: IInlineChatRequest = {
447-
prompt: this._activeSession.lastInput,
458+
prompt: this._activeSession.lastInput.value,
459+
attempt: this._activeSession.lastInput.attempt,
448460
selection: this._editor.getSelection(),
449461
wholeRange: this._activeSession.wholeRange.value,
450-
attempt: 0,
451462
};
452463
const task = this._activeSession.provider.provideResponse(this._activeSession.session, request, requestCts.token);
453464
this._log('request started', this._activeSession.provider.debugName, this._activeSession.session, request);
@@ -463,7 +474,7 @@ export class InlineChatController implements IEditorContribution {
463474
if (reply?.type === 'message') {
464475
response = new MarkdownResponse(this._activeSession.textModelN.uri, reply);
465476
} else if (reply) {
466-
response = new EditResponse(this._activeSession.textModelN.uri, reply);
477+
response = new EditResponse(this._activeSession.textModelN.uri, this._activeSession.textModelN.getAlternativeVersionId(), reply);
467478
} else {
468479
response = new EmptyResponse();
469480
}
@@ -483,7 +494,7 @@ export class InlineChatController implements IEditorContribution {
483494
msgListener.dispose();
484495
typeListener.dispose();
485496

486-
this._activeSession.addExchange(new SessionExchange(request.prompt, response));
497+
this._activeSession.addExchange(new SessionExchange(this._activeSession.lastInput, response));
487498

488499
if (message & Message.CANCEL_SESSION) {
489500
return State.CANCEL;
@@ -646,6 +657,10 @@ export class InlineChatController implements IEditorContribution {
646657
this._messages.fire(Message.ACCEPT_INPUT);
647658
}
648659

660+
regenerate(): void {
661+
this._messages.fire(Message.RERUN_INPUT);
662+
}
663+
649664
cancelCurrentRequest(): void {
650665
this._messages.fire(Message.CANCEL_INPUT | Message.CANCEL_REQUEST);
651666
}
@@ -683,7 +698,7 @@ export class InlineChatController implements IEditorContribution {
683698

684699
viewInChat() {
685700
if (this._activeSession?.lastExchange?.response instanceof MarkdownResponse) {
686-
this._instaService.invokeFunction(showMessageResponse, this._activeSession.lastExchange.prompt, this._activeSession.lastExchange.response.raw.message.value);
701+
this._instaService.invokeFunction(showMessageResponse, this._activeSession.lastExchange.prompt.value, this._activeSession.lastExchange.response.raw.message.value);
687702
}
688703
}
689704

‎src/vs/workbench/contrib/inlineChat/browser/inlineChatLivePreviewWidget.ts

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ export class InlineChatLivePreviewWidget extends ZoneWidget {
4444
private readonly _inlineDiffDecorations: IEditorDecorationsCollection;
4545
private _dim: Dimension | undefined;
4646
private _isVisible: boolean = false;
47+
private _isDiffLocked: boolean = false;
4748

4849
constructor(
4950
editor: ICodeEditor,
@@ -134,6 +135,8 @@ export class InlineChatLivePreviewWidget extends ZoneWidget {
134135
override show(): void {
135136
assertType(this.editor.hasModel());
136137
this._sessionStore.clear();
138+
this._isDiffLocked = false;
139+
this._isVisible = true;
137140

138141
this._sessionStore.add(this._diffEditor.onDidUpdateDiff(() => {
139142
const result = this._diffEditor.getDiffComputationResult();
@@ -148,12 +151,19 @@ export class InlineChatLivePreviewWidget extends ZoneWidget {
148151
}
149152
}));
150153
this._updateFromChanges(this._session.wholeRange.value, this._session.lastTextModelChanges);
151-
this._isVisible = true;
154+
}
155+
156+
lockToDiff(): void {
157+
this._isDiffLocked = true;
152158
}
153159

154160
private _updateFromChanges(range: Range, changes: readonly LineRangeMapping[]): void {
155161
assertType(this.editor.hasModel());
156162

163+
if (this._isDiffLocked) {
164+
return;
165+
}
166+
157167
if (changes.length === 0 || this._session.textModel0.getValueLength() === 0) {
158168
// no change or changes to an empty file
159169
this._logService.debug('[IE] livePreview-mode: no diff');

‎src/vs/workbench/contrib/inlineChat/browser/inlineChatSession.ts

Lines changed: 33 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,7 @@ class SessionWholeRange {
108108

109109
export class Session {
110110

111-
private _lastInput: string | undefined;
111+
private _lastInput: SessionPrompt | undefined;
112112
private _lastExpansionState: ExpansionState | undefined;
113113
private _lastTextModelChanges: readonly LineRangeMapping[] | undefined;
114114
private _isUnstashed: boolean = false;
@@ -139,10 +139,14 @@ export class Session {
139139
};
140140
}
141141

142-
addInput(input: string): void {
142+
addInput(input: SessionPrompt): void {
143143
this._lastInput = input;
144144
}
145145

146+
get lastInput() {
147+
return this._lastInput;
148+
}
149+
146150
get isUnstashed(): boolean {
147151
return this._isUnstashed;
148152
}
@@ -151,10 +155,6 @@ export class Session {
151155
this._isUnstashed = true;
152156
}
153157

154-
get lastInput() {
155-
return this._lastInput;
156-
}
157-
158158
get lastExpansionState(): ExpansionState | undefined {
159159
return this._lastExpansionState;
160160
}
@@ -229,17 +229,37 @@ export class Session {
229229
for (const exchange of this._exchange) {
230230
const response = exchange.response;
231231
if (response instanceof MarkdownResponse || response instanceof EditResponse) {
232-
result.exchanges.push({ prompt: exchange.prompt, res: response.raw });
232+
result.exchanges.push({ prompt: exchange.prompt.value, res: response.raw });
233233
}
234234
}
235235
return result;
236236
}
237237
}
238238

239239

240+
export class SessionPrompt {
241+
242+
private _attempt: number = 0;
243+
244+
constructor(
245+
readonly value: string,
246+
) { }
247+
248+
get attempt() {
249+
return this._attempt;
250+
}
251+
252+
retry() {
253+
const result = new SessionPrompt(this.value);
254+
result._attempt = this._attempt + 1;
255+
return result;
256+
}
257+
}
258+
240259
export class SessionExchange {
260+
241261
constructor(
242-
readonly prompt: string,
262+
readonly prompt: SessionPrompt,
243263
readonly response: MarkdownResponse | EditResponse | EmptyResponse | ErrorResponse
244264
) { }
245265
}
@@ -275,7 +295,11 @@ export class EditResponse {
275295
readonly workspaceEdits: ResourceEdit[] | undefined;
276296
readonly workspaceEditsIncludeLocalEdits: boolean = false;
277297

278-
constructor(localUri: URI, readonly raw: IInlineChatBulkEditResponse | IInlineChatEditResponse) {
298+
constructor(
299+
localUri: URI,
300+
readonly modelAltVersionId: number,
301+
readonly raw: IInlineChatBulkEditResponse | IInlineChatEditResponse
302+
) {
279303
if (raw.type === 'editorEdit') {
280304
//
281305
this.localEdits = raw.edits;

‎src/vs/workbench/contrib/inlineChat/browser/inlineChatStrategies.ts

Lines changed: 24 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ import { Position } from 'vs/editor/common/core/position';
1313
import { Range } from 'vs/editor/common/core/range';
1414
import { Selection } from 'vs/editor/common/core/selection';
1515
import { IEditorDecorationsCollection } from 'vs/editor/common/editorCommon';
16-
import { ICursorStateComputer, IModelDecorationOptions, IModelDeltaDecoration, IValidEditOperation } from 'vs/editor/common/model';
16+
import { ICursorStateComputer, IModelDecorationOptions, IModelDeltaDecoration, ITextModel, IValidEditOperation } from 'vs/editor/common/model';
1717
import { IEditorWorkerService } from 'vs/editor/common/services/editorWorker';
1818
import { localize } from 'vs/nls';
1919
import { IContextKey, IContextKeyService } from 'vs/platform/contextkey/common/contextkey';
@@ -37,6 +37,8 @@ export abstract class EditModeStrategy {
3737

3838
abstract makeChanges(edits: ISingleEditOperation[]): Promise<void>;
3939

40+
abstract undoChanges(response: EditResponse): Promise<void>;
41+
4042
abstract renderChanges(response: EditResponse): Promise<void>;
4143

4244
abstract toggleDiff(): void;
@@ -110,6 +112,10 @@ export class PreviewStrategy extends EditModeStrategy {
110112
// nothing to do
111113
}
112114

115+
override async undoChanges(_response: EditResponse): Promise<void> {
116+
// nothing to do
117+
}
118+
113119
override async renderChanges(response: EditResponse): Promise<void> {
114120
if (response.localEdits.length > 0) {
115121
const edits = response.localEdits.map(edit => EditOperation.replace(Range.lift(edit.range), edit.text));
@@ -275,9 +281,7 @@ export class LiveStrategy extends EditModeStrategy {
275281
return;
276282
}
277283
const targetAltVersion = textModelNSnapshotAltVersion ?? textModelNAltVersion;
278-
while (targetAltVersion < modelN.getAlternativeVersionId() && modelN.canUndo()) {
279-
modelN.undo();
280-
}
284+
LiveStrategy._undoModelUntil(modelN, targetAltVersion);
281285
}
282286

283287
override async makeChanges(edits: ISingleEditOperation[], ignoreInlineDiff?: boolean): Promise<void> {
@@ -297,6 +301,11 @@ export class LiveStrategy extends EditModeStrategy {
297301
this._editor.executeEdits('inline-chat-live', edits, ignoreInlineDiff ? undefined : cursorStateComputerAndInlineDiffCollection);
298302
}
299303

304+
override async undoChanges(response: EditResponse): Promise<void> {
305+
const { textModelN } = this._session;
306+
LiveStrategy._undoModelUntil(textModelN, response.modelAltVersionId);
307+
}
308+
300309
override async renderChanges(response: EditResponse) {
301310

302311
this._inlineDiffDecorations.update();
@@ -309,6 +318,12 @@ export class LiveStrategy extends EditModeStrategy {
309318
}
310319
}
311320

321+
private static _undoModelUntil(model: ITextModel, targetAltVersion: number): void {
322+
while (targetAltVersion < model.getAlternativeVersionId() && model.canUndo()) {
323+
model.undo();
324+
}
325+
}
326+
312327
protected _updateSummaryMessage() {
313328
let linesChanged = 0;
314329
for (const change of this._session.lastTextModelChanges) {
@@ -373,6 +388,11 @@ export class LivePreviewStrategy extends LiveStrategy {
373388
}
374389
}
375390

391+
override async undoChanges(response: EditResponse): Promise<void> {
392+
this._diffZone.lockToDiff();
393+
super.undoChanges(response);
394+
}
395+
376396
protected override _doToggleDiff(): void {
377397
const scrollState = StableEditorScrollState.capture(this._editor);
378398
if (this._diffEnabled) {

0 commit comments

Comments
 (0)