Skip to content

Commit f0d1b49

Browse files
jriekenCopilot
andauthored
Inline chat handles empty selections explicitly (#2535)
* handle empty selection better in inline chat * stests * add unit tests * Update src/extension/prompts/node/inline/inlineChat2Prompt.tsx Co-authored-by: Copilot <[email protected]> * stests --------- Co-authored-by: Copilot <[email protected]>
1 parent 5b32ec3 commit f0d1b49

13 files changed

+804
-442
lines changed

src/extension/prompts/node/inline/inlineChat2Prompt.tsx

Lines changed: 84 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,12 @@
44
*--------------------------------------------------------------------------------------------*/
55

66
import { AssistantMessage, PromptElement, PromptElementProps, PromptReference, PromptSizing, SystemMessage, ToolCall, ToolMessage, useKeepWith, UserMessage } from '@vscode/prompt-tsx';
7-
import type { ExtendedLanguageModelToolResult } from 'vscode';
7+
import { ChatResponsePart } from '@vscode/prompt-tsx/dist/base/vscodeTypes';
8+
import type { CancellationToken, ExtendedLanguageModelToolResult, Position, Progress } from 'vscode';
89
import { TextDocumentSnapshot } from '../../../../platform/editing/common/textDocumentSnapshot';
910
import { CacheType } from '../../../../platform/endpoint/common/endpointTypes';
1011
import { IPromptPathRepresentationService } from '../../../../platform/prompts/common/promptPathRepresentationService';
11-
import { ChatRequest, ChatRequestEditorData } from '../../../../vscodeTypes';
12+
import { ChatRequest, ChatRequestEditorData, Range } from '../../../../vscodeTypes';
1213
import { ChatVariablesCollection } from '../../../prompt/common/chatVariablesCollection';
1314
import { IToolCall } from '../../../prompt/common/intents';
1415
import { CopilotIdentityRules } from '../base/copilotIdentity';
@@ -41,15 +42,7 @@ export class InlineChat2Prompt extends PromptElement<InlineChat2PromptProps> {
4142

4243
const snapshotAtRequest = this.props.snapshotAtRequest;
4344

44-
// the full lines of the selection
45-
// TODO@jrieken
46-
// * if the selection is empty and if the line with the selection is empty we could hint to add code and
47-
// generally with empty selections we could allow the model to be a bit more creative
48-
// * use the true selected text (now we extend to full lines)
49-
const selectedLines = snapshotAtRequest.getText(this.props.data.selection.with({
50-
start: this.props.data.selection.start.with({ character: 0 }),
51-
end: this.props.data.selection.end.with({ character: Number.MAX_SAFE_INTEGER }),
52-
}));
45+
const selection = this.props.data.selection;
5346

5447
const variables = new ChatVariablesCollection(this.props.request.references);
5548
const filepath = this._promptPathRepresentationService.getFilePath(snapshotAtRequest.uri);
@@ -74,9 +67,10 @@ export class InlineChat2Prompt extends PromptElement<InlineChat2PromptProps> {
7467
<Tag name='file'>
7568
<CodeBlock includeFilepath={false} languageId={snapshotAtRequest.languageId} uri={snapshotAtRequest.uri} references={[new PromptReference(snapshotAtRequest.uri, undefined, undefined)]} code={snapshotAtRequest.getText()} />
7669
</Tag>
77-
<Tag name='file-selection'>
78-
<CodeBlock includeFilepath={false} languageId={snapshotAtRequest.languageId} uri={snapshotAtRequest.uri} references={[new PromptReference(snapshotAtRequest.uri, undefined, undefined)]} code={selectedLines} />
79-
</Tag>
70+
{selection.isEmpty
71+
? <FileContextElement snapshot={snapshotAtRequest} position={selection.start} />
72+
: <FileSelectionElement snapshot={snapshotAtRequest} selection={selection} />
73+
}
8074
<ChatVariables flexGrow={3} priority={898} chatVariables={variables} useFixCookbook={true} />
8175
<Tag name='reminder'>
8276
If there is a user selection, focus on it, and try to make changes to the selected code and its context.<br />
@@ -98,7 +92,82 @@ export class InlineChat2Prompt extends PromptElement<InlineChat2PromptProps> {
9892
}
9993
}
10094

101-
export type EditAttemptsElementProps = PromptElementProps<{
95+
96+
export type FileContextElementProps = PromptElementProps<{
97+
snapshot: TextDocumentSnapshot;
98+
position: Position;
99+
}>;
100+
101+
export class FileContextElement extends PromptElement<FileContextElementProps> {
102+
103+
override render(state: void, sizing: PromptSizing, _progress?: Progress<ChatResponsePart>, _token?: CancellationToken) {
104+
105+
let startLine = this.props.position.line;
106+
let endLine = this.props.position.line;
107+
let n = 0;
108+
let seenNonEmpty = false;
109+
while (startLine > 0) {
110+
seenNonEmpty = seenNonEmpty || !this.props.snapshot.lineAt(startLine).isEmptyOrWhitespace;
111+
startLine--;
112+
n++;
113+
if (n >= 3 && seenNonEmpty) {
114+
break;
115+
}
116+
}
117+
n = 0;
118+
seenNonEmpty = false;
119+
while (endLine < this.props.snapshot.lineCount - 1) {
120+
seenNonEmpty = seenNonEmpty || !this.props.snapshot.lineAt(endLine).isEmptyOrWhitespace;
121+
endLine++;
122+
n++;
123+
if (n >= 3 && seenNonEmpty) {
124+
break;
125+
}
126+
}
127+
128+
const textBefore = this.props.snapshot.getText(new Range(this.props.position.with({ line: startLine, character: 0 }), this.props.position));
129+
const textAfter = this.props.snapshot.getText(new Range(this.props.position, this.props.position.with({ line: endLine, character: Number.MAX_SAFE_INTEGER })));
130+
131+
const code = `${textBefore}$CURSOR$${textAfter}`;
132+
133+
return <>
134+
<Tag name='file-cursor-context'>
135+
<CodeBlock includeFilepath={false} languageId={this.props.snapshot.languageId} uri={this.props.snapshot.uri} references={[new PromptReference(this.props.snapshot.uri, undefined, undefined)]} code={code} />
136+
</Tag>
137+
</>;
138+
}
139+
}
140+
141+
142+
export type FileSelectionElementProps = PromptElementProps<{
143+
snapshot: TextDocumentSnapshot;
144+
selection: Range;
145+
}>;
146+
147+
export class FileSelectionElement extends PromptElement<FileSelectionElementProps> {
148+
149+
override render(state: void, sizing: PromptSizing, progress?: Progress<ChatResponsePart>, token?: CancellationToken) {
150+
151+
152+
// the full lines of the selection
153+
// TODO@jrieken
154+
// * use the true selected text (now we extend to full lines)
155+
156+
const selectedLines = this.props.snapshot.getText(this.props.selection.with({
157+
start: this.props.selection.start.with({ character: 0 }),
158+
end: this.props.selection.end.with({ character: Number.MAX_SAFE_INTEGER }),
159+
}));
160+
161+
return <>
162+
<Tag name='file-selection'>
163+
<CodeBlock includeFilepath={false} languageId={this.props.snapshot.languageId} uri={this.props.snapshot.uri} references={[new PromptReference(this.props.snapshot.uri, undefined, undefined)]} code={selectedLines} />
164+
</Tag>
165+
</>;
166+
}
167+
}
168+
169+
170+
type EditAttemptsElementProps = PromptElementProps<{
102171
editAttempts: [IToolCall, ExtendedLanguageModelToolResult][];
103172
data: ChatRequestEditorData;
104173
documentVersionAtRequest: number;
Lines changed: 243 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,243 @@
1+
/*---------------------------------------------------------------------------------------------
2+
* Copyright (c) Microsoft Corporation. All rights reserved.
3+
* Licensed under the MIT License. See License.txt in the project root for license information.
4+
*--------------------------------------------------------------------------------------------*/
5+
6+
import { expect, suite, test } from 'vitest';
7+
import { TextDocumentSnapshot } from '../../../../../platform/editing/common/textDocumentSnapshot';
8+
import { createTextDocumentData } from '../../../../../util/common/test/shims/textDocument';
9+
import { URI } from '../../../../../util/vs/base/common/uri';
10+
import { Position, Range } from '../../../../../vscodeTypes';
11+
import { FileContextElement, FileSelectionElement } from '../inlineChat2Prompt';
12+
13+
function createSnapshot(content: string, languageId: string = 'typescript'): TextDocumentSnapshot {
14+
const uri = URI.file('/workspace/file.ts');
15+
const docData = createTextDocumentData(uri, content, languageId);
16+
return TextDocumentSnapshot.create(docData.document);
17+
}
18+
19+
suite('FileContextElement', () => {
20+
21+
test('cursor at the beginning of the file', async () => {
22+
const content = `line 1
23+
line 2
24+
line 3
25+
line 4
26+
line 5`;
27+
const snapshot = createSnapshot(content);
28+
const position = new Position(0, 0);
29+
30+
const element = new FileContextElement({ snapshot, position });
31+
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });
32+
33+
const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
34+
expect(output).toContain('$CURSOR$');
35+
expect(output).toContain('line 1');
36+
expect(output).toContain('line 2');
37+
expect(output).toContain('line 3');
38+
});
39+
40+
test('cursor in the middle of a file', async () => {
41+
const content = `line 1
42+
line 2
43+
line 3
44+
line 4
45+
line 5
46+
line 6
47+
line 7`;
48+
const snapshot = createSnapshot(content);
49+
const position = new Position(3, 2); // after "li" in "line 4"
50+
51+
const element = new FileContextElement({ snapshot, position });
52+
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });
53+
54+
const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
55+
expect(output).toContain('$CURSOR$');
56+
// Should include lines before and after cursor
57+
expect(output).toContain('line 2');
58+
expect(output).toContain('line 3');
59+
// Cursor position (3, 2) splits "line 4" into "li" + "$CURSOR$" + "ne 4"
60+
expect(output).toContain('li$CURSOR$ne 4');
61+
expect(output).toContain('line 5');
62+
expect(output).toContain('line 6');
63+
});
64+
65+
test('cursor at the end of file', async () => {
66+
const content = `line 1
67+
line 2
68+
line 3
69+
line 4
70+
line 5`;
71+
const snapshot = createSnapshot(content);
72+
const position = new Position(4, 6); // end of "line 5"
73+
74+
const element = new FileContextElement({ snapshot, position });
75+
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });
76+
77+
const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
78+
expect(output).toContain('$CURSOR$');
79+
expect(output).toContain('line 3');
80+
expect(output).toContain('line 4');
81+
expect(output).toContain('line 5');
82+
});
83+
84+
test('cursor with empty lines - includes extra lines until non-empty', async () => {
85+
const content = `
86+
87+
line 3
88+
line 4
89+
90+
`;
91+
const snapshot = createSnapshot(content);
92+
const position = new Position(2, 0); // start of "line 3"
93+
94+
const element = new FileContextElement({ snapshot, position });
95+
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });
96+
97+
const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
98+
expect(output).toContain('$CURSOR$');
99+
expect(output).toContain('line 3');
100+
expect(output).toContain('line 4');
101+
});
102+
103+
test('single line file', async () => {
104+
const content = `only one line`;
105+
const snapshot = createSnapshot(content);
106+
const position = new Position(0, 5); // middle of line
107+
108+
const element = new FileContextElement({ snapshot, position });
109+
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });
110+
111+
const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
112+
expect(output).toContain('only $CURSOR$one line');
113+
});
114+
115+
test('cursor position splits text correctly', async () => {
116+
const content = `hello world`;
117+
const snapshot = createSnapshot(content);
118+
const position = new Position(0, 6); // after "hello "
119+
120+
const element = new FileContextElement({ snapshot, position });
121+
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });
122+
123+
const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
124+
expect(output).toContain('hello $CURSOR$world');
125+
});
126+
});
127+
128+
suite('FileSelectionElement', () => {
129+
130+
test('single line selection', async () => {
131+
const content = `line 1
132+
line 2
133+
line 3
134+
line 4
135+
line 5`;
136+
const snapshot = createSnapshot(content);
137+
const selection = new Range(1, 0, 1, 6); // "line 2"
138+
139+
const element = new FileSelectionElement({ snapshot, selection });
140+
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });
141+
142+
const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
143+
expect(output).toContain('line 2');
144+
expect(output).not.toContain('line 1');
145+
expect(output).not.toContain('line 3');
146+
});
147+
148+
test('multi-line selection', async () => {
149+
const content = `line 1
150+
line 2
151+
line 3
152+
line 4
153+
line 5`;
154+
const snapshot = createSnapshot(content);
155+
const selection = new Range(1, 0, 3, 6); // "line 2" through "line 4"
156+
157+
const element = new FileSelectionElement({ snapshot, selection });
158+
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });
159+
160+
const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
161+
expect(output).toContain('line 2');
162+
expect(output).toContain('line 3');
163+
expect(output).toContain('line 4');
164+
expect(output).not.toContain('line 1');
165+
expect(output).not.toContain('line 5');
166+
});
167+
168+
test('partial line selection extends to full lines', async () => {
169+
const content = `line 1
170+
line 2
171+
line 3`;
172+
const snapshot = createSnapshot(content);
173+
// Select from middle of line 2 to middle of line 2 (partial)
174+
const selection = new Range(1, 2, 1, 4);
175+
176+
const element = new FileSelectionElement({ snapshot, selection });
177+
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });
178+
179+
const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
180+
// Should include the full line, not just "ne"
181+
expect(output).toContain('line 2');
182+
});
183+
184+
test('selection at start of file', async () => {
185+
const content = `line 1
186+
line 2
187+
line 3`;
188+
const snapshot = createSnapshot(content);
189+
const selection = new Range(0, 0, 0, 6);
190+
191+
const element = new FileSelectionElement({ snapshot, selection });
192+
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });
193+
194+
const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
195+
expect(output).toContain('line 1');
196+
expect(output).not.toContain('line 2');
197+
});
198+
199+
test('selection at end of file', async () => {
200+
const content = `line 1
201+
line 2
202+
line 3`;
203+
const snapshot = createSnapshot(content);
204+
const selection = new Range(2, 0, 2, 6);
205+
206+
const element = new FileSelectionElement({ snapshot, selection });
207+
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });
208+
209+
const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
210+
expect(output).toContain('line 3');
211+
expect(output).not.toContain('line 2');
212+
});
213+
214+
test('selection spanning partial lines extends to full lines', async () => {
215+
const content = `first line here
216+
second line here
217+
third line here`;
218+
const snapshot = createSnapshot(content);
219+
// Select from middle of "first" to middle of "second"
220+
const selection = new Range(0, 6, 1, 7);
221+
222+
const element = new FileSelectionElement({ snapshot, selection });
223+
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });
224+
225+
const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
226+
// Should include full lines
227+
expect(output).toContain('first line here');
228+
expect(output).toContain('second line here');
229+
expect(output).not.toContain('third line here');
230+
});
231+
232+
test('preserves language id for code block', async () => {
233+
const content = `const x = 1;`;
234+
const snapshot = createSnapshot(content, 'javascript');
235+
const selection = new Range(0, 0, 0, 12);
236+
237+
const element = new FileSelectionElement({ snapshot, selection });
238+
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });
239+
240+
const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
241+
expect(output).toContain('javascript');
242+
});
243+
});

0 commit comments

Comments
 (0)