Skip to content

Commit 8725d13

Browse files
haritabh-z01haritabh-z01
authored
chore: migrate ipcs to trpc for main window (#24)
* chore: migrate ipcs to trpc for main window * chore: switch to trpc subscriptions for model downloads
1 parent d7481f7 commit 8725d13

File tree

10 files changed

+649
-698
lines changed

10 files changed

+649
-698
lines changed

‎apps/electron/src/components/models-view.tsx

Lines changed: 152 additions & 122 deletions
Original file line numberDiff line numberDiff line change
@@ -19,148 +19,173 @@ import {
1919
AlertDialogCancel
2020
} from './ui/alert-dialog';
2121
import { Model, DownloadedModel, DownloadProgress } from '../constants/models';
22+
import { api } from '@/trpc/react';
2223

2324
export const ModelsView: React.FC = () => {
24-
const [availableModels, setAvailableModels] = useState<Model[]>([]);
25-
const [downloadedModels, setDownloadedModels] = useState<Record<string, DownloadedModel>>({});
2625
const [downloadProgress, setDownloadProgress] = useState<Record<string, DownloadProgress>>({});
27-
const [loading, setLoading] = useState(true);
28-
const [isLocalWhisperAvailable, setIsLocalWhisperAvailable] = useState(false);
29-
const [selectedModel, setSelectedModel] = useState<string | null>(null);
3026
const [showDeleteDialog, setShowDeleteDialog] = useState(false);
3127
const [modelToDelete, setModelToDelete] = useState<string | null>(null);
3228

33-
const loadData = async () => {
34-
try {
35-
setLoading(true);
36-
const [available, downloaded, activeDownloads, whisperAvailable, currentSelectedModel] = await Promise.all([
37-
window.electronAPI.getAvailableModels(),
38-
window.electronAPI.getDownloadedModels(),
39-
window.electronAPI.getActiveDownloads(),
40-
window.electronAPI.isLocalWhisperAvailable(),
41-
window.electronAPI.getSelectedModel(),
42-
]);
43-
44-
setAvailableModels(available);
45-
setDownloadedModels(downloaded);
46-
setIsLocalWhisperAvailable(whisperAvailable);
47-
setSelectedModel(currentSelectedModel);
48-
49-
// Set up active downloads progress
50-
const progressMap: Record<string, DownloadProgress> = {};
51-
for (const downloadProgress of activeDownloads) {
52-
progressMap[downloadProgress.modelId] = downloadProgress;
29+
// tRPC queries
30+
const availableModelsQuery = api.models.getAvailableModels.useQuery();
31+
const downloadedModelsQuery = api.models.getDownloadedModels.useQuery();
32+
const activeDownloadsQuery = api.models.getActiveDownloads.useQuery();
33+
const isLocalWhisperAvailableQuery = api.models.isLocalWhisperAvailable.useQuery();
34+
const selectedModelQuery = api.models.getSelectedModel.useQuery();
35+
36+
const utils = api.useUtils();
37+
38+
// tRPC mutations
39+
const downloadModelMutation = api.models.downloadModel.useMutation({
40+
onSuccess: () => {
41+
utils.models.getDownloadedModels.invalidate();
42+
utils.models.getActiveDownloads.invalidate();
43+
},
44+
onError: (error) => {
45+
console.error('Failed to start download:', error);
46+
if (error instanceof Error && error.message.includes('AbortError')) {
47+
console.log('Download was manually aborted, not showing error');
48+
return;
5349
}
54-
setDownloadProgress(progressMap);
55-
} catch (err) {
56-
console.error('Failed to load models data:', err);
57-
toast.error('Failed to load models data');
58-
} finally {
59-
setLoading(false);
50+
toast.error('Failed to start download');
6051
}
61-
};
52+
});
6253

54+
const cancelDownloadMutation = api.models.cancelDownload.useMutation({
55+
onSuccess: () => {
56+
utils.models.getActiveDownloads.invalidate();
57+
},
58+
onError: (error) => {
59+
console.error('Failed to cancel download:', error);
60+
toast.error('Failed to cancel download');
61+
}
62+
});
63+
64+
const deleteModelMutation = api.models.deleteModel.useMutation({
65+
onSuccess: () => {
66+
utils.models.getDownloadedModels.invalidate();
67+
setShowDeleteDialog(false);
68+
setModelToDelete(null);
69+
},
70+
onError: (error) => {
71+
console.error('Failed to delete model:', error);
72+
toast.error('Failed to delete model');
73+
setShowDeleteDialog(false);
74+
setModelToDelete(null);
75+
}
76+
});
77+
78+
const setSelectedModelMutation = api.models.setSelectedModel.useMutation({
79+
onSuccess: () => {
80+
utils.models.getSelectedModel.invalidate();
81+
},
82+
onError: (error) => {
83+
console.error('Failed to select model:', error);
84+
toast.error('Failed to select model');
85+
}
86+
});
87+
88+
// Initialize active downloads progress on load
6389
useEffect(() => {
64-
loadData();
65-
66-
const handleDownloadProgress = (modelId: string, progress: DownloadProgress) => {
67-
setDownloadProgress(prev => ({
68-
...prev,
69-
[modelId]: progress
70-
}));
71-
};
72-
73-
const handleDownloadComplete = (modelId: string, downloadedModel: DownloadedModel) => {
74-
setDownloadedModels(prev => ({
75-
...prev,
76-
[modelId]: downloadedModel
77-
}));
90+
if (activeDownloadsQuery.data) {
91+
const progressMap: Record<string, DownloadProgress> = {};
92+
activeDownloadsQuery.data.forEach((download) => {
93+
progressMap[download.modelId] = download;
94+
});
95+
setDownloadProgress(progressMap);
96+
}
97+
}, [activeDownloadsQuery.data]);
98+
99+
// Set up tRPC subscriptions for real-time download updates
100+
api.models.onDownloadProgress.useSubscription(undefined, {
101+
onData: ({ modelId, progress }) => {
102+
setDownloadProgress(prev => ({ ...prev, [modelId]: progress }));
103+
},
104+
onError: (error) => {
105+
console.error('Download progress subscription error:', error);
106+
}
107+
});
108+
109+
api.models.onDownloadComplete.useSubscription(undefined, {
110+
onData: ({ modelId, downloadedModel }) => {
78111
setDownloadProgress(prev => {
79-
const updated = { ...prev };
80-
delete updated[modelId];
81-
return updated;
112+
const newProgress = { ...prev };
113+
delete newProgress[modelId];
114+
return newProgress;
82115
});
83-
};
84-
85-
const handleDownloadError = (modelId: string, errorMessage: string) => {
86-
setDownloadProgress(prev => ({
87-
...prev,
88-
[modelId]: {
89-
...prev[modelId],
90-
status: 'error',
91-
error: errorMessage
92-
}
93-
}));
94-
};
95-
96-
const handleDownloadCancelled = (modelId: string) => {
116+
utils.models.getDownloadedModels.invalidate();
117+
utils.models.getActiveDownloads.invalidate();
118+
},
119+
onError: (error) => {
120+
console.error('Download complete subscription error:', error);
121+
}
122+
});
123+
124+
api.models.onDownloadError.useSubscription(undefined, {
125+
onData: ({ modelId, error }) => {
97126
setDownloadProgress(prev => {
98-
const updated = { ...prev };
99-
delete updated[modelId];
100-
return updated;
127+
const newProgress = { ...prev };
128+
delete newProgress[modelId];
129+
return newProgress;
101130
});
102-
};
131+
toast.error(`Download failed: ${error}`);
132+
utils.models.getActiveDownloads.invalidate();
133+
},
134+
onError: (error) => {
135+
console.error('Download error subscription error:', error);
136+
}
137+
});
103138

104-
const handleModelDeleted = (modelId: string) => {
105-
setDownloadedModels(prev => {
106-
const updated = { ...prev };
107-
delete updated[modelId];
108-
return updated;
139+
api.models.onDownloadCancelled.useSubscription(undefined, {
140+
onData: ({ modelId }) => {
141+
setDownloadProgress(prev => {
142+
const newProgress = { ...prev };
143+
delete newProgress[modelId];
144+
return newProgress;
109145
});
110-
};
111-
112-
// Listen to events from main process
113-
window.electronAPI.on('model-download-progress', handleDownloadProgress);
114-
window.electronAPI.on('model-download-complete', handleDownloadComplete);
115-
window.electronAPI.on('model-download-error', handleDownloadError);
116-
window.electronAPI.on('model-download-cancelled', handleDownloadCancelled);
117-
window.electronAPI.on('model-deleted', handleModelDeleted);
118-
119-
return () => {
120-
// Cleanup event listeners
121-
window.electronAPI.off('model-download-progress', handleDownloadProgress);
122-
window.electronAPI.off('model-download-complete', handleDownloadComplete);
123-
window.electronAPI.off('model-download-error', handleDownloadError);
124-
window.electronAPI.off('model-download-cancelled', handleDownloadCancelled);
125-
window.electronAPI.off('model-deleted', handleModelDeleted);
126-
};
127-
}, []);
146+
utils.models.getActiveDownloads.invalidate();
147+
},
148+
onError: (error) => {
149+
console.error('Download cancelled subscription error:', error);
150+
}
151+
});
152+
153+
api.models.onModelDeleted.useSubscription(undefined, {
154+
onData: ({ modelId }) => {
155+
utils.models.getDownloadedModels.invalidate();
156+
},
157+
onError: (error) => {
158+
console.error('Model deleted subscription error:', error);
159+
}
160+
});
128161

129162
const handleDownload = async (modelId: string, event?: React.MouseEvent) => {
130-
event?.preventDefault();
131-
event?.stopPropagation();
132-
133-
console.log('Start download clicked for:', modelId);
134-
163+
if (event) {
164+
event.preventDefault();
165+
event.stopPropagation();
166+
}
167+
135168
try {
136-
console.log('Downloading model:', modelId);
137-
await window.electronAPI.downloadModel(modelId);
138-
console.log('Start download successful for:', modelId);
169+
await downloadModelMutation.mutateAsync({ modelId });
170+
console.log('Download started for:', modelId);
139171
} catch (err) {
140172
console.error('Failed to start download:', err);
141-
142-
// Don't show error for manual cancellations (AbortError)
143-
if (err instanceof Error && err.message.includes('AbortError')) {
144-
console.log('Download was manually aborted, not showing error');
145-
return;
146-
}
147-
148-
toast.error('Failed to start download');
173+
// Error is already handled by the mutation's onError
149174
}
150175
};
151176

152177
const handleCancelDownload = async (modelId: string, event?: React.MouseEvent) => {
153-
event?.preventDefault();
154-
event?.stopPropagation();
155-
156-
console.log('Cancel download clicked for:', modelId);
157-
178+
if (event) {
179+
event.preventDefault();
180+
event.stopPropagation();
181+
}
182+
158183
try {
159-
await window.electronAPI.cancelDownload(modelId);
184+
await cancelDownloadMutation.mutateAsync({ modelId });
160185
console.log('Cancel download successful for:', modelId);
161186
} catch (err) {
162187
console.error('Failed to cancel download:', err);
163-
toast.error('Failed to cancel download');
188+
// Error is already handled by the mutation's onError
164189
}
165190
};
166191

@@ -173,13 +198,10 @@ export const ModelsView: React.FC = () => {
173198
if (!modelToDelete) return;
174199

175200
try {
176-
await window.electronAPI.deleteModel(modelToDelete);
201+
await deleteModelMutation.mutateAsync({ modelId: modelToDelete });
177202
} catch (err) {
178203
console.error('Failed to delete model:', err);
179-
toast.error('Failed to delete model');
180-
} finally {
181-
setShowDeleteDialog(false);
182-
setModelToDelete(null);
204+
// Error is already handled by the mutation's onError
183205
}
184206
};
185207

@@ -188,14 +210,12 @@ export const ModelsView: React.FC = () => {
188210
setModelToDelete(null);
189211
};
190212

191-
192213
const handleSelectModel = async (modelId: string) => {
193214
try {
194-
await window.electronAPI.setSelectedModel(modelId);
195-
setSelectedModel(modelId);
215+
await setSelectedModelMutation.mutateAsync({ modelId });
196216
} catch (err) {
197217
console.error('Failed to select model:', err);
198-
toast.error('Failed to select model');
218+
// Error is already handled by the mutation's onError
199219
}
200220
};
201221

@@ -207,6 +227,16 @@ export const ModelsView: React.FC = () => {
207227
return parseFloat((bytes / Math.pow(k, i)).toFixed(1)) + ' ' + sizes[i];
208228
};
209229

230+
// Loading state
231+
const loading = availableModelsQuery.isLoading || downloadedModelsQuery.isLoading ||
232+
isLocalWhisperAvailableQuery.isLoading || selectedModelQuery.isLoading;
233+
234+
// Data from queries
235+
const availableModels = availableModelsQuery.data || [];
236+
const downloadedModels = downloadedModelsQuery.data || {};
237+
const isLocalWhisperAvailable = isLocalWhisperAvailableQuery.data || false;
238+
const selectedModel = selectedModelQuery.data;
239+
210240
if (loading) {
211241
return (
212242
<div className="flex items-center justify-center h-64">

0 commit comments

Comments
 (0)