store.ts

src/state/store.ts

1. import { create } from 'zustand'
2. import { useShallow } from 'zustand/react/shallow'
4. import type { IpcClient } from '../ipc/client'
5. import type { AppInfo } from '../ipc/gen/AppInfo'
6. import type { AttachmentMeta } from '../ipc/gen/AttachmentMeta'
7. import type { ConversationSummary } from '../ipc/gen/ConversationSummary'
8. import type { ConversationView } from '../ipc/gen/ConversationView'
9. import type { Message } from '../ipc/gen/Message'
10. import type { ModelInfo } from '../ipc/gen/ModelInfo'
11. import type { ModelsResponse } from '../ipc/gen/ModelsResponse'
12. import type { Preset } from '../ipc/gen/Preset'
13. import type { Provider } from '../ipc/gen/Provider'
14. import type { Settings } from '../ipc/gen/Settings'
15. import type { SpendSummary } from '../ipc/gen/SpendSummary'
16. import type { StreamEvent } from '../ipc/gen/StreamEvent'
18. export type Panel = null | 'search' | 'settings' | 'usage'
20. export interface Streaming {
21. streamId: string
22. assistantId: string
23. text: string
24. reasoning: string
25. }
27. export interface Toast {
28. id: number
29. text: string
30. tone: 'info' | 'error'
31. }
33. interface State {
34. client: IpcClient | null
35. info: AppInfo | null
36. settings: Settings | null
37. providers: Provider[]
38. models: Record<string, ModelsResponse>
39. keys: Record<string, boolean>
40. presets: Preset[]
41. conversations: ConversationSummary[]
42. view: ConversationView | null
43. streaming: Streaming | null
44. drafts: AttachmentMeta[]
45. spend: SpendSummary | null
46. panel: Panel
47. toasts: Toast[]
48. /** For a new chat not yet created: the model it will use. */
49. pending: { providerId: string; modelId: string; presetId: string | null }
51. init(client: IpcClient): Promise<void>
52. toast(text: string, tone?: Toast['tone']): void
53. setPanel(p: Panel): void
54. refreshConversations(): Promise<void>
55. refreshSpend(): Promise<void>
56. refreshModels(providerId: string, live?: boolean): Promise<void>
57. refreshProviders(): Promise<void>
58. open(id: string): Promise<void>
59. newChat(): void
60. currentModel(): { providerId: string; modelId: string; info: ModelInfo | null }
61. chooseModel(providerId: string, modelId: string): Promise<void>
62. applyPreset(preset: Preset | null): Promise<void>
63. send(text: string): Promise<boolean>
64. stop(): Promise<void>
65. regenerate(): Promise<void>
66. editAndResend(messageId: string, text: string): Promise<void>
67. addDrafts(a: AttachmentMeta[]): void
68. removeDraft(id: string): Promise<void>
69. saveSettings(s: Settings): Promise<void>
70. }
72. let toastSeq = 0
73. // Deltas pile up here and reach React once per frame.
74. let pendingText = ''
75. let pendingReasoning = ''
76. let frame = 0
78. export const useStore = create<State>((set, get) => {
79. const c = () => {
80. const client = get().client
81. if (!client) throw new Error('not initialized')
82. return client
83. }
85. const flush = () => {
86. frame = 0
87. const s = get().streaming
88. if (!s || (!pendingText && !pendingReasoning)) return
89. set({ streaming: { ...s, text: s.text + pendingText, reasoning: s.reasoning + pendingReasoning } })
90. pendingText = ''
91. pendingReasoning = ''
92. }
93. const schedule = () => {
94. if (frame) return
95. frame =
96. typeof requestAnimationFrame === 'function'
97. ? requestAnimationFrame(flush)
98. : (setTimeout(flush, 16) as unknown as number)
99. }
101. const finish = (message: Message) => {
102. flush()
103. const v = get().view
104. if (v && v.conversation.id === message.conversation_id) {
105. set({ view: { ...v, messages: v.messages.map((m) => (m.id === message.id ? message : m)) } })
106. }
107. set({ streaming: null })
108. void get().refreshConversations()
109. void get().refreshSpend()
110. }
112. const onEvent = (e: StreamEvent) => {
113. switch (e.type) {
114. case 'delta':
115. pendingText += e.text
116. schedule()
117. break
118. case 'reasoning':
119. pendingReasoning += e.text
120. schedule()
121. break
122. case 'usage':
123. break
124. case 'done':
125. finish(e.message)
126. break
127. case 'error':
128. finish(e.message)
129. get().toast(e.error, 'error')
130. break
131. }
132. }
134. const startTurn = async (conversationId: string, text: string, attachmentIds: string[]) => {
135. pendingText = ''
136. pendingReasoning = ''
137. const ack = await c().sendMessage({ conversation_id: conversationId, text, attachment_ids: attachmentIds }, onEvent)
138. const v = get().view
139. if (v && v.conversation.id === conversationId) {
140. const add = [ack.user_message, ack.assistant_message].filter((m): m is Message => !!m)
141. set({ view: { ...v, messages: [...v.messages, ...add] } })
142. }
143. // A fast stream can finish before the ack; only mark it live if it has not.
144. const live = get().view?.messages.find((m) => m.id === ack.assistant_message.id)
145. if (!live || live.status === 'streaming') {
146. set({
147. streaming: {
148. streamId: ack.stream_id,
149. assistantId: ack.assistant_message.id,
150. text: pendingText,
151. reasoning: pendingReasoning,
152. },
153. })
154. pendingText = ''
155. pendingReasoning = ''
156. }
157. }
159. return {
160. client: null,
161. info: null,
162. settings: null,
163. providers: [],
164. models: {},
165. keys: {},
166. presets: [],
167. conversations: [],
168. view: null,
169. streaming: null,
170. drafts: [],
171. spend: null,
172. panel: null,
173. toasts: [],
174. pending: { providerId: 'deepinfra', modelId: '', presetId: null },
176. async init(client) {
177. set({ client })
178. // Listen first: the shell refreshes model prices at startup and may
179. // announce it before the first reads below have finished.
180. await Promise.all([
181. client.on('conversations-changed', () => void get().refreshConversations()),
182. client.on('models-updated', () => get().providers.forEach((p) => void get().refreshModels(p.id))),
183. client.on('settings-changed', (s) => set({ settings: s as Settings })),
184. client.on('open-conversation', (id) => void get().open(id as string)),
185. client.on('new-chat', () => get().newChat()),
186. client.on('open-settings', () => set({ panel: 'settings' })),
187. ])
188. const [info, settings, providers, presets] = await Promise.all([
189. client.appInfo(),
190. client.getSettings(),
191. client.listProviders(),
192. client.listPresets(),
193. ])
194. set({
195. info,
196. settings,
197. providers,
198. presets,
199. pending: { providerId: settings.default_provider_id, modelId: settings.default_model_id, presetId: null },
200. })
201. await Promise.all([
202. ...providers.map((p) => get().refreshModels(p.id)),
203. get().refreshProviders(),
204. get().refreshConversations(),
205. get().refreshSpend(),
206. ])
207. },
209. toast(text, tone = 'info') {
210. const id = ++toastSeq
211. set({ toasts: [...get().toasts, { id, text, tone }] })
212. setTimeout(() => set({ toasts: get().toasts.filter((t) => t.id !== id) }), tone === 'error' ? 7000 : 3500)
213. },
215. setPanel(panel) {
216. set({ panel })
217. },
219. async refreshConversations() {
220. const conversations = await c().listConversations(false)
221. set({ conversations })
222. // The chat is titled from its first message on the Rust side; keep the
223. // open one in step, or its header says "New chat" until reopened.
224. const v = get().view
225. const mine = v && conversations.find((x) => x.id === v.conversation.id)
226. if (v && mine && (mine.title !== v.conversation.title || mine.pinned !== v.conversation.pinned)) {
227. set({ view: { ...v, conversation: { ...v.conversation, title: mine.title, pinned: mine.pinned } } })
228. }
229. },
231. async refreshSpend() {
232. set({ spend: await c().spendSummary() })
233. },
235. async refreshModels(providerId, live = false) {
236. const r = live ? await c().refreshModels(providerId) : await c().listModels(providerId)
237. set({ models: { ...get().models, [providerId]: r } })
238. if (live && r.warning) get().toast(`Could not refresh models: ${r.warning}`, 'error')
239. },
241. async refreshProviders() {
242. const providers = await c().listProviders()
243. const keys: Record<string, boolean> = {}
244. await Promise.all(providers.map(async (p) => (keys[p.id] = await c().hasApiKey(p.id))))
245. set({ providers, keys })
246. },
248. async open(id) {
249. const view = await c().getConversation(id)
250. set({ view, panel: null })
251. },
253. newChat() {
254. const s = get().settings
255. const cur = get().currentModel()
256. set({
257. view: null,
258. drafts: [],
259. panel: null,
260. pending: {
261. providerId: cur.providerId || s?.default_provider_id || 'deepinfra',
262. modelId: cur.modelId || s?.default_model_id || '',
263. presetId: null,
264. },
265. })
266. },
268. currentModel() {
269. const v = get().view
270. const providerId = v ? v.conversation.provider_id : get().pending.providerId
271. const modelId = v ? v.conversation.model_id : get().pending.modelId
272. const info = get().models[providerId]?.models.find((m) => m.id === modelId) ?? null
273. return { providerId, modelId, info }
274. },
276. async chooseModel(providerId, modelId) {
277. const v = get().view
278. if (v) {
279. await c().setConversationModel(v.conversation.id, providerId, modelId)
280. set({ view: { ...v, conversation: { ...v.conversation, provider_id: providerId, model_id: modelId } } })
281. } else {
282. set({ pending: { ...get().pending, providerId, modelId } })
283. }
284. },
286. async applyPreset(preset) {
287. const v = get().view
288. if (v) {
289. await c().setSystemPrompt(v.conversation.id, preset?.system_prompt ?? '', preset?.id ?? null)
290. if (preset?.provider_id && preset.model_id)
291. await c().setConversationModel(v.conversation.id, preset.provider_id, preset.model_id)
292. await get().open(v.conversation.id)
293. } else {
294. set({
295. pending: {
296. providerId: preset?.provider_id ?? get().pending.providerId,
297. modelId: preset?.model_id ?? get().pending.modelId,
298. presetId: preset?.id ?? null,
299. },
300. })
301. }
302. },
304. async send(text) {
305. if (get().streaming) return false
306. const drafts = get().drafts
307. if (!text.trim() && !drafts.length) return false
308. const s = get().settings
309. const spend = get().spend
310. if (
311. s?.confirm_over_budget &&
312. spend?.budget_level === 2 &&
313. !window.confirm("You are over this month's budget. Send anyway?")
314. )
315. return false
316. try {
317. let v = get().view
318. if (!v) {
319. const p = get().pending
320. const conv = await c().createConversation(p.providerId, p.modelId, p.presetId, 'main')
321. v = { conversation: conv, messages: [] }
322. set({ view: v })
323. }
324. set({ drafts: [] })
325. await startTurn(
326. v.conversation.id,
327. text,
328. drafts.map((d) => d.id),
329. )
330. void get().refreshConversations()
331. return true
332. } catch (err) {
333. set({ drafts })
334. get().toast(String(err), 'error')
335. return false
336. }
337. },
339. async stop() {
340. const s = get().streaming
341. if (s) await c().cancelStream(s.streamId)
342. },
344. async regenerate() {
345. const v = get().view
346. if (!v || get().streaming) return
347. const last = [...v.messages].reverse().find((m) => m.role === 'assistant')
348. if (!last) return
349. await c().truncateFrom(v.conversation.id, last.id)
350. await get().open(v.conversation.id)
351. try {
352. await startTurn(v.conversation.id, '', [])
353. } catch (err) {
354. get().toast(String(err), 'error')
355. }
356. },
358. async editAndResend(messageId, text) {
359. const v = get().view
360. if (!v || get().streaming) return
361. const original = v.messages.find((m) => m.id === messageId)
362. await c().truncateFrom(v.conversation.id, messageId)
363. await get().open(v.conversation.id)
364. try {
365. await startTurn(v.conversation.id, text, original?.attachments.map((a) => a.id) ?? [])
366. } catch (err) {
367. get().toast(String(err), 'error')
368. }
369. },
371. addDrafts(a) {
372. set({ drafts: [...get().drafts, ...a] })
373. },
375. async removeDraft(id) {
376. set({ drafts: get().drafts.filter((d) => d.id !== id) })
377. await c().discardAttachment(id)
378. },
380. async saveSettings(s) {
381. const r = await c().saveSettings(s)
382. set({ settings: r.settings })
383. if (r.shortcut_error) get().toast(r.shortcut_error, 'error')
384. void get().refreshSpend()
385. },
386. }
387. })
389. /** The model the current (or next) chat uses. Compared field by field, so it
390. * only re-renders when the model actually changes. */
391. export const useCurrentModel = () => useStore(useShallow((s) => s.currentModel()))