Skip to content

Commit 9b1bdc4

Browse files
robertpiosikrobertpiosik
committed
Optimize token calculation progress notification threshold, synchronize selected files on initialization, and debounce selected files component state updates
- should show 'crunching...' notification if the prcess takes longer than 3 seconds - on initial load we get token count but not amount of files - the `packages/ui/src/components/editor/prompt/SelectedFiles/SelectedFiles.tsx` shows updated file and token count in a slight desynchronization ├── apps │ └── editor │ └── src │ ├── context │ │ ├── context-initialization.ts │ │ ├── helpers │ │ │ ├── get-all-workspace-files.ts │ │ │ └── get-selected-files.ts │ │ ├── providers │ │ │ ├── open-editors │ │ │ │ └── open-editors-provider.ts │ │ │ ├── selected-files │ │ │ │ └── selected-files-provider.ts │ │ │ └── workspace │ │ │ ├── modules │ │ │ │ ├── index.ts │ │ │ │ └── token-calculator.ts │ │ │ └── workspace-provider.ts │ │ ├── shared-context-state.ts │ │ └── utils │ │ └── shrink-file │ │ ├── index.ts │ │ ├── languages │ │ │ ├── c-style.ts │ │ │ ├── css.ts │ │ │ ├── html.ts │ │ │ ├── jsx.ts │ │ │ ├── python.ts │ │ │ ├── ruby.ts │ │ │ └── sql.ts │ │ ├── shrink-file.spec.ts │ │ ├── shrink-file.ts │ │ └── test-cases │ │ ├── c-style │ │ │ ├── input.txt │ │ │ └── output.txt │ │ ├── css │ │ │ ├── input.txt │ │ │ └── output.txt │ │ ├── html │ │ │ ├── input.txt │ │ │ └── output.txt │ │ ├── jsx │ │ │ ├── input.txt │ │ │ └── output.txt │ │ ├── python │ │ │ ├── input.txt │ │ │ └── output.txt │ │ ├── ruby │ │ │ ├── input.txt │ │ │ └── output.txt │ │ ├── scss │ │ │ ├── input.txt │ │ │ └── output.txt │ │ └── sql │ │ ├── input.txt │ │ └── output.txt │ ├── i18n │ │ └── translations │ │ └── context.ts │ └── views │ └── prompt │ ├── backend │ │ └── prompt-view-provider.ts │ ├── frontend │ │ ├── Main │ │ │ ├── Main.tsx │ │ │ └── MainView │ │ │ └── MainView.tsx │ │ ├── Prompt.tsx │ │ └── hooks │ │ └── use-prompt │ │ └── use-prompt.ts │ └── types │ └── messages.ts └── packages └── ui └── src └── components └── editor └── prompt └── SelectedFiles └── SelectedFiles.tsx
1 parent 7d456cc commit 9b1bdc4

5 files changed

Lines changed: 95 additions & 30 deletions

File tree

‎apps/editor/src/context/providers/workspace/modules/token-calculator.ts‎

Lines changed: 70 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ export class TokenCalculator implements vscode.Disposable {
4040
private _notified_roots: Set<string> = new Set()
4141

4242
private _progress_stack: {
43-
progress: vscode.Progress<{ message?: string; increment?: number }>
43+
progress?: vscode.Progress<{ message?: string; increment?: number }>
4444
total_files: number
4545
processed_files: number
4646
}[] = []
@@ -171,7 +171,9 @@ export class TokenCalculator implements vscode.Disposable {
171171
if (state.total_files > 0) {
172172
state.processed_files++
173173
const increment = (1 / state.total_files) * 100
174-
state.progress.report({ increment })
174+
if (state.progress) {
175+
state.progress.report({ increment })
176+
}
175177
}
176178
}
177179
}
@@ -274,30 +276,75 @@ export class TokenCalculator implements vscode.Disposable {
274276
return params.task()
275277
}
276278

277-
return vscode.window.withProgress(
278-
{
279-
location: vscode.ProgressLocation.Notification,
280-
title: t('context.crunching-token-counts'),
281-
cancellable: false
282-
},
283-
async (progress) => {
284-
const progress_state = {
285-
progress,
286-
total_files: total_uncached,
287-
processed_files: 0
288-
}
289-
this._progress_stack.push(progress_state)
290-
291-
try {
292-
return await params.task()
293-
} finally {
294-
const index = this._progress_stack.indexOf(progress_state)
295-
if (index !== -1) {
296-
this._progress_stack.splice(index, 1)
279+
const progress_state: {
280+
progress?: vscode.Progress<{ message?: string; increment?: number }>
281+
total_files: number
282+
processed_files: number
283+
} = {
284+
total_files: total_uncached,
285+
processed_files: 0
286+
}
287+
this._progress_stack.push(progress_state)
288+
289+
const task_promise = params.task()
290+
291+
let timeout_handle: NodeJS.Timeout
292+
const timeout_promise = new Promise<void>((resolve) => {
293+
timeout_handle = setTimeout(() => resolve(), 3000)
294+
})
295+
296+
let is_completed = false
297+
task_promise.finally(() => {
298+
is_completed = true
299+
clearTimeout(timeout_handle)
300+
})
301+
302+
try {
303+
await Promise.race([task_promise, timeout_promise])
304+
} catch (error) {
305+
const index = this._progress_stack.indexOf(progress_state)
306+
if (index !== -1) {
307+
this._progress_stack.splice(index, 1)
308+
}
309+
throw error
310+
}
311+
312+
if (!is_completed) {
313+
return vscode.window.withProgress(
314+
{
315+
location: vscode.ProgressLocation.Notification,
316+
title: t('context.crunching-token-counts'),
317+
cancellable: false
318+
},
319+
async (progress) => {
320+
progress_state.progress = progress
321+
if (
322+
progress_state.total_files > 0 &&
323+
progress_state.processed_files > 0
324+
) {
325+
progress.report({
326+
increment:
327+
(progress_state.processed_files / progress_state.total_files) *
328+
100
329+
})
330+
}
331+
try {
332+
return await task_promise
333+
} finally {
334+
const index = this._progress_stack.indexOf(progress_state)
335+
if (index !== -1) {
336+
this._progress_stack.splice(index, 1)
337+
}
297338
}
298339
}
340+
)
341+
} else {
342+
const index = this._progress_stack.indexOf(progress_state)
343+
if (index !== -1) {
344+
this._progress_stack.splice(index, 1)
299345
}
300-
)
346+
return task_promise
347+
}
301348
}
302349

303350
public invalidate_token_counts_for_file(changed_file_path: string) {

‎apps/editor/src/views/prompt/backend/prompt-view-provider.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -854,6 +854,8 @@ export class PromptViewProvider implements vscode.WebviewViewProvider {
854854
await handle_pick_tasks_workspace(this, message)
855855
} else if (message.command == 'GET_TOKEN_COUNT') {
856856
this.send_token_count()
857+
} else if (message.command == 'GET_SELECTED_FILES') {
858+
this.send_selected_files()
857859
} else if (message.command == 'PREVIEW_PROMPT') {
858860
await handle_preview_prompt({
859861
prompt_view_provider: this

‎apps/editor/src/views/prompt/frontend/hooks/use-prompt/use-prompt.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -150,7 +150,8 @@ export const use_prompt = (vscode: any) => {
150150
{ command: 'GET_SETUP_PROGRESS' },
151151
{ command: 'GET_VOICE_INPUT_PUSH_TO_TALK' },
152152
{ command: 'GET_IS_MODERN_UI' },
153-
{ command: 'GET_TOKEN_COUNT' }
153+
{ command: 'GET_TOKEN_COUNT' },
154+
{ command: 'GET_SELECTED_FILES' }
154155
]
155156
initial_messages.forEach((message) => post_message(vscode, message))
156157

‎apps/editor/src/views/prompt/types/messages.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -508,6 +508,10 @@ export interface GetTokenCountMessage extends BaseMessage {
508508
command: 'GET_TOKEN_COUNT'
509509
}
510510

511+
export interface GetSelectedFilesMessage extends BaseMessage {
512+
command: 'GET_SELECTED_FILES'
513+
}
514+
511515
export interface PreviewPromptMessage extends BaseMessage {
512516
command: 'PREVIEW_PROMPT'
513517
}
@@ -600,6 +604,7 @@ export type FrontendMessage =
600604
| PickTasksWorkspaceMessage
601605
| GetIsModernUiMessage
602606
| GetTokenCountMessage
607+
| GetSelectedFilesMessage
603608
| PreviewPromptMessage
604609

605610
// === FROM BACKEND TO FRONTEND ===

‎packages/ui/src/components/editor/prompt/SelectedFiles/SelectedFiles.tsx‎

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,21 @@ type Props = {
1111
}
1212

1313
export const SelectedFiles: React.FC<Props> = (props) => {
14-
const [files_ount, set_files_count] = useState(props.files_count)
14+
const [displayed_state, set_displayed_state] = useState({
15+
token_count: props.token_count,
16+
files_count: props.files_count
17+
})
1518

1619
useEffect(() => {
17-
set_files_count(props.files_count)
18-
}, [props.token_count])
20+
const timeout = setTimeout(() => {
21+
set_displayed_state({
22+
token_count: props.token_count,
23+
files_count: props.files_count
24+
})
25+
}, 50)
26+
27+
return () => clearTimeout(timeout)
28+
}, [props.token_count, props.files_count])
1929

2030
return (
2131
<div className={styles.container}>
@@ -26,12 +36,12 @@ export const SelectedFiles: React.FC<Props> = (props) => {
2636
<span>
2737
{props.translations.attaching_files.replace(
2838
'{files}',
29-
String(files_ount)
39+
String(displayed_state.files_count)
3040
)}
3141
</span>
32-
{props.token_count > 0 && (
42+
{displayed_state.token_count > 0 && (
3343
<span className={styles.tokens}>
34-
{display_token_count(props.token_count)}
44+
{display_token_count(displayed_state.token_count)}
3545
</span>
3646
)}
3747
</label>

0 commit comments

Comments
 (0)