1
1
import { useLoaderData , useNavigate , useSearchParams } from '@remix-run/react' ;
2
- import { useState , useEffect } from 'react' ;
2
+ import { useState , useEffect , useCallback } from 'react' ;
3
3
import { atom } from 'nanostores' ;
4
- import type { Message } from 'ai' ;
4
+ import { generateId , type JSONValue , type Message } from 'ai' ;
5
5
import { toast } from 'react-toastify' ;
6
6
import { workbenchStore } from '~/lib/stores/workbench' ;
7
7
import { logStore } from '~/lib/stores/logs' ; // Import logStore
@@ -15,6 +15,11 @@ import {
15
15
createChatFromMessages ,
16
16
type IChatMetadata ,
17
17
} from './db' ;
18
+ import type { FileMap } from '~/lib/stores/files' ;
19
+ import type { Snapshot } from './types' ;
20
+ import { webcontainer } from '~/lib/webcontainer' ;
21
+ import { createCommandsMessage , detectProjectCommands } from '~/utils/projectCommands' ;
22
+ import type { ContextAnnotation } from '~/types/context' ;
18
23
19
24
export interface ChatHistoryItem {
20
25
id : string ;
@@ -37,6 +42,7 @@ export function useChatHistory() {
37
42
const { id : mixedId } = useLoaderData < { id ?: string } > ( ) ;
38
43
const [ searchParams ] = useSearchParams ( ) ;
39
44
45
+ const [ archivedMessages , setArchivedMessages ] = useState < Message [ ] > ( [ ] ) ;
40
46
const [ initialMessages , setInitialMessages ] = useState < Message [ ] > ( [ ] ) ;
41
47
const [ ready , setReady ] = useState < boolean > ( false ) ;
42
48
const [ urlId , setUrlId ] = useState < string | undefined > ( ) ;
@@ -56,14 +62,128 @@ export function useChatHistory() {
56
62
57
63
if ( mixedId ) {
58
64
getMessages ( db , mixedId )
59
- . then ( ( storedMessages ) => {
65
+ . then ( async ( storedMessages ) => {
60
66
if ( storedMessages && storedMessages . messages . length > 0 ) {
67
+ const snapshotStr = localStorage . getItem ( `snapshot:${ mixedId } ` ) ;
68
+ const snapshot : Snapshot = snapshotStr ? JSON . parse ( snapshotStr ) : { chatIndex : 0 , files : { } } ;
69
+ const summary = snapshot . summary ;
70
+
61
71
const rewindId = searchParams . get ( 'rewindTo' ) ;
62
- const filteredMessages = rewindId
63
- ? storedMessages . messages . slice ( 0 , storedMessages . messages . findIndex ( ( m ) => m . id === rewindId ) + 1 )
64
- : storedMessages . messages ;
72
+ let startingIdx = - 1 ;
73
+ const endingIdx = rewindId
74
+ ? storedMessages . messages . findIndex ( ( m ) => m . id === rewindId ) + 1
75
+ : storedMessages . messages . length ;
76
+ const snapshotIndex = storedMessages . messages . findIndex ( ( m ) => m . id === snapshot . chatIndex ) ;
77
+
78
+ if ( snapshotIndex >= 0 && snapshotIndex < endingIdx ) {
79
+ startingIdx = snapshotIndex ;
80
+ }
81
+
82
+ if ( snapshotIndex > 0 && storedMessages . messages [ snapshotIndex ] . id == rewindId ) {
83
+ startingIdx = - 1 ;
84
+ }
85
+
86
+ let filteredMessages = storedMessages . messages . slice ( startingIdx + 1 , endingIdx ) ;
87
+ let archivedMessages : Message [ ] = [ ] ;
88
+
89
+ if ( startingIdx >= 0 ) {
90
+ archivedMessages = storedMessages . messages . slice ( 0 , startingIdx + 1 ) ;
91
+ }
92
+
93
+ setArchivedMessages ( archivedMessages ) ;
94
+
95
+ if ( startingIdx > 0 ) {
96
+ const files = Object . entries ( snapshot ?. files || { } )
97
+ . map ( ( [ key , value ] ) => {
98
+ if ( value ?. type !== 'file' ) {
99
+ return null ;
100
+ }
101
+
102
+ return {
103
+ content : value . content ,
104
+ path : key ,
105
+ } ;
106
+ } )
107
+ . filter ( ( x ) => ! ! x ) ;
108
+ const projectCommands = await detectProjectCommands ( files ) ;
109
+ const commands = createCommandsMessage ( projectCommands ) ;
110
+
111
+ filteredMessages = [
112
+ {
113
+ id : generateId ( ) ,
114
+ role : 'user' ,
115
+ content : `Restore project from snapshot
116
+ ` ,
117
+ annotations : [ 'no-store' , 'hidden' ] ,
118
+ } ,
119
+ {
120
+ id : storedMessages . messages [ snapshotIndex ] . id ,
121
+ role : 'assistant' ,
122
+ content : ` 📦 Chat Restored from snapshot, You can revert this message to load the full chat history
123
+ <boltArtifact id="imported-files" title="Project Files Snapshot" type="bundled">
124
+ ${ Object . entries ( snapshot ?. files || { } )
125
+ . filter ( ( x ) => ! x [ 0 ] . endsWith ( 'lock.json' ) )
126
+ . map ( ( [ key , value ] ) => {
127
+ if ( value ?. type === 'file' ) {
128
+ return `
129
+ <boltAction type="file" filePath="${ key } ">
130
+ ${ value . content }
131
+ </boltAction>
132
+ ` ;
133
+ } else {
134
+ return `` ;
135
+ }
136
+ } )
137
+ . join ( '\n' ) }
138
+ </boltArtifact>
139
+ ` ,
140
+ annotations : [
141
+ 'no-store' ,
142
+ ...( summary
143
+ ? [
144
+ {
145
+ chatId : storedMessages . messages [ snapshotIndex ] . id ,
146
+ type : 'chatSummary' ,
147
+ summary,
148
+ } satisfies ContextAnnotation ,
149
+ ]
150
+ : [ ] ) ,
151
+ ] ,
152
+ } ,
153
+ ...( commands !== null
154
+ ? [
155
+ {
156
+ id : `${ storedMessages . messages [ snapshotIndex ] . id } -2` ,
157
+ role : 'user' as const ,
158
+ content : `setup project` ,
159
+ annotations : [ 'no-store' , 'hidden' ] ,
160
+ } ,
161
+ {
162
+ ...commands ,
163
+ id : `${ storedMessages . messages [ snapshotIndex ] . id } -3` ,
164
+ annotations : [
165
+ 'no-store' ,
166
+ ...( commands . annotations || [ ] ) ,
167
+ ...( summary
168
+ ? [
169
+ {
170
+ chatId : `${ storedMessages . messages [ snapshotIndex ] . id } -3` ,
171
+ type : 'chatSummary' ,
172
+ summary,
173
+ } satisfies ContextAnnotation ,
174
+ ]
175
+ : [ ] ) ,
176
+ ] ,
177
+ } ,
178
+ ]
179
+ : [ ] ) ,
180
+ ...filteredMessages ,
181
+ ] ;
182
+ restoreSnapshot ( mixedId ) ;
183
+ }
65
184
66
185
setInitialMessages ( filteredMessages ) ;
186
+
67
187
setUrlId ( storedMessages . urlId ) ;
68
188
description . set ( storedMessages . description ) ;
69
189
chatId . set ( storedMessages . id ) ;
@@ -75,10 +195,64 @@ export function useChatHistory() {
75
195
setReady ( true ) ;
76
196
} )
77
197
. catch ( ( error ) => {
198
+ console . error ( error ) ;
199
+
78
200
logStore . logError ( 'Failed to load chat messages' , error ) ;
79
201
toast . error ( error . message ) ;
80
202
} ) ;
81
203
}
204
+ } , [ mixedId ] ) ;
205
+
206
+ const takeSnapshot = useCallback (
207
+ async ( chatIdx : string , files : FileMap , _chatId ?: string | undefined , chatSummary ?: string ) => {
208
+ const id = _chatId || chatId ;
209
+
210
+ if ( ! id ) {
211
+ return ;
212
+ }
213
+
214
+ const snapshot : Snapshot = {
215
+ chatIndex : chatIdx ,
216
+ files,
217
+ summary : chatSummary ,
218
+ } ;
219
+ localStorage . setItem ( `snapshot:${ id } ` , JSON . stringify ( snapshot ) ) ;
220
+ } ,
221
+ [ chatId ] ,
222
+ ) ;
223
+
224
+ const restoreSnapshot = useCallback ( async ( id : string ) => {
225
+ const snapshotStr = localStorage . getItem ( `snapshot:${ id } ` ) ;
226
+ const container = await webcontainer ;
227
+
228
+ // if (snapshotStr)setSnapshot(JSON.parse(snapshotStr));
229
+ const snapshot : Snapshot = snapshotStr ? JSON . parse ( snapshotStr ) : { chatIndex : 0 , files : { } } ;
230
+
231
+ if ( ! snapshot ?. files ) {
232
+ return ;
233
+ }
234
+
235
+ Object . entries ( snapshot . files ) . forEach ( async ( [ key , value ] ) => {
236
+ if ( key . startsWith ( container . workdir ) ) {
237
+ key = key . replace ( container . workdir , '' ) ;
238
+ }
239
+
240
+ if ( value ?. type === 'folder' ) {
241
+ await container . fs . mkdir ( key , { recursive : true } ) ;
242
+ }
243
+ } ) ;
244
+ Object . entries ( snapshot . files ) . forEach ( async ( [ key , value ] ) => {
245
+ if ( value ?. type === 'file' ) {
246
+ if ( key . startsWith ( container . workdir ) ) {
247
+ key = key . replace ( container . workdir , '' ) ;
248
+ }
249
+
250
+ await container . fs . writeFile ( key , value . content , { encoding : value . isBinary ? undefined : 'utf8' } ) ;
251
+ } else {
252
+ }
253
+ } ) ;
254
+
255
+ // workbenchStore.files.setKey(snapshot?.files)
82
256
} , [ ] ) ;
83
257
84
258
return {
@@ -105,14 +279,34 @@ export function useChatHistory() {
105
279
}
106
280
107
281
const { firstArtifact } = workbenchStore ;
282
+ messages = messages . filter ( ( m ) => ! m . annotations ?. includes ( 'no-store' ) ) ;
283
+
284
+ let _urlId = urlId ;
108
285
109
286
if ( ! urlId && firstArtifact ?. id ) {
110
287
const urlId = await getUrlId ( db , firstArtifact . id ) ;
111
-
288
+ _urlId = urlId ;
112
289
navigateChat ( urlId ) ;
113
290
setUrlId ( urlId ) ;
114
291
}
115
292
293
+ let chatSummary : string | undefined = undefined ;
294
+ const lastMessage = messages [ messages . length - 1 ] ;
295
+
296
+ if ( lastMessage . role === 'assistant' ) {
297
+ const annotations = lastMessage . annotations as JSONValue [ ] ;
298
+ const filteredAnnotations = ( annotations ?. filter (
299
+ ( annotation : JSONValue ) =>
300
+ annotation && typeof annotation === 'object' && Object . keys ( annotation ) . includes ( 'type' ) ,
301
+ ) || [ ] ) as { type : string ; value : any } & { [ key : string ] : any } [ ] ;
302
+
303
+ if ( filteredAnnotations . find ( ( annotation ) => annotation . type === 'chatSummary' ) ) {
304
+ chatSummary = filteredAnnotations . find ( ( annotation ) => annotation . type === 'chatSummary' ) ?. summary ;
305
+ }
306
+ }
307
+
308
+ takeSnapshot ( messages [ messages . length - 1 ] . id , workbenchStore . files . get ( ) , _urlId , chatSummary ) ;
309
+
116
310
if ( ! description . get ( ) && firstArtifact ?. title ) {
117
311
description . set ( firstArtifact ?. title ) ;
118
312
}
@@ -127,7 +321,15 @@ export function useChatHistory() {
127
321
}
128
322
}
129
323
130
- await setMessages ( db , chatId . get ( ) as string , messages , urlId , description . get ( ) , undefined , chatMetadata . get ( ) ) ;
324
+ await setMessages (
325
+ db ,
326
+ chatId . get ( ) as string ,
327
+ [ ...archivedMessages , ...messages ] ,
328
+ urlId ,
329
+ description . get ( ) ,
330
+ undefined ,
331
+ chatMetadata . get ( ) ,
332
+ ) ;
131
333
} ,
132
334
duplicateCurrentChat : async ( listItemId : string ) => {
133
335
if ( ! db || ( ! mixedId && ! listItemId ) ) {
0 commit comments