@@ -307,8 +307,30 @@ class ChatStore {
307307 onError ?: ( error : Error ) => void
308308 ) : Promise < void > {
309309 let streamedContent = '' ;
310-
311310 let streamedReasoningContent = '' ;
311+ let modelCaptured = false ;
312+
313+ const captureModelIfNeeded = ( updateDbImmediately = true ) : string | undefined => {
314+ if ( ! modelCaptured ) {
315+ const currentModelName = serverStore . modelName ;
316+
317+ if ( currentModelName ) {
318+ if ( updateDbImmediately ) {
319+ DatabaseStore . updateMessage ( assistantMessage . id , { model : currentModelName } ) . catch (
320+ console . error
321+ ) ;
322+ }
323+
324+ const messageIndex = this . findMessageIndex ( assistantMessage . id ) ;
325+
326+ this . updateMessageAtIndex ( messageIndex , { model : currentModelName } ) ;
327+ modelCaptured = true ;
328+
329+ return currentModelName ;
330+ }
331+ }
332+ return undefined ;
333+ } ;
312334
313335 slotsService . startStreaming ( ) ;
314336
@@ -319,6 +341,8 @@ class ChatStore {
319341 streamedContent += chunk ;
320342 this . currentResponse = streamedContent ;
321343
344+ captureModelIfNeeded ( ) ;
345+
322346 const partialThinking = extractPartialThinking ( streamedContent ) ;
323347 const messageIndex = this . findMessageIndex ( assistantMessage . id ) ;
324348 this . updateMessageAtIndex ( messageIndex , {
@@ -328,7 +352,11 @@ class ChatStore {
328352
329353 onReasoningChunk : ( reasoningChunk : string ) => {
330354 streamedReasoningContent += reasoningChunk ;
355+
356+ captureModelIfNeeded ( ) ;
357+
331358 const messageIndex = this . findMessageIndex ( assistantMessage . id ) ;
359+
332360 this . updateMessageAtIndex ( messageIndex , { thinking : streamedReasoningContent } ) ;
333361 } ,
334362
@@ -339,17 +367,36 @@ class ChatStore {
339367 ) => {
340368 slotsService . stopStreaming ( ) ;
341369
342- await DatabaseStore . updateMessage ( assistantMessage . id , {
370+ const updateData : {
371+ content : string ;
372+ thinking : string ;
373+ timings ?: ChatMessageTimings ;
374+ model ?: string ;
375+ } = {
343376 content : finalContent || streamedContent ,
344377 thinking : reasoningContent || streamedReasoningContent ,
345378 timings : timings
346- } ) ;
379+ } ;
380+
381+ const capturedModel = captureModelIfNeeded ( false ) ;
382+
383+ if ( capturedModel ) {
384+ updateData . model = capturedModel ;
385+ }
386+
387+ await DatabaseStore . updateMessage ( assistantMessage . id , updateData ) ;
347388
348389 const messageIndex = this . findMessageIndex ( assistantMessage . id ) ;
349390
350- this . updateMessageAtIndex ( messageIndex , {
391+ const localUpdateData : { timings ?: ChatMessageTimings ; model ?: string } = {
351392 timings : timings
352- } ) ;
393+ } ;
394+
395+ if ( updateData . model ) {
396+ localUpdateData . model = updateData . model ;
397+ }
398+
399+ this . updateMessageAtIndex ( messageIndex , localUpdateData ) ;
353400
354401 await DatabaseStore . updateCurrentNode ( this . activeConversation ! . id , assistantMessage . id ) ;
355402 this . activeConversation ! . currNode = assistantMessage . id ;
@@ -478,9 +525,6 @@ class ChatStore {
478525 private async createAssistantMessage ( parentId ?: string ) : Promise < DatabaseMessage | null > {
479526 if ( ! this . activeConversation ) return null ;
480527
481- // Capture the current model name when creating the assistant message
482- const currentModelName = serverStore . modelName ;
483-
484528 return await DatabaseStore . createMessageBranch (
485529 {
486530 convId : this . activeConversation . id ,
@@ -489,8 +533,7 @@ class ChatStore {
489533 content : '' ,
490534 timestamp : Date . now ( ) ,
491535 thinking : '' ,
492- children : [ ] ,
493- model : currentModelName || undefined
536+ children : [ ]
494537 } ,
495538 parentId || null
496539 ) ;
@@ -1287,9 +1330,6 @@ class ChatStore {
12871330 this . isLoading = true ;
12881331 this . currentResponse = '' ;
12891332
1290- // Capture the current model name when creating the assistant message
1291- const currentModelName = serverStore . modelName ;
1292-
12931333 const newAssistantMessage = await DatabaseStore . createMessageBranch (
12941334 {
12951335 convId : this . activeConversation . id ,
@@ -1298,8 +1338,7 @@ class ChatStore {
12981338 role : 'assistant' ,
12991339 content : '' ,
13001340 thinking : '' ,
1301- children : [ ] ,
1302- model : currentModelName || undefined
1341+ children : [ ]
13031342 } ,
13041343 parentMessage . id
13051344 ) ;
@@ -1346,9 +1385,6 @@ class ChatStore {
13461385 false
13471386 ) as DatabaseMessage [ ] ;
13481387
1349- // Capture the current model name when creating the assistant message
1350- const currentModelName = serverStore . modelName ;
1351-
13521388 // Create new assistant message branch
13531389 const assistantMessage = await DatabaseStore . createMessageBranch (
13541390 {
@@ -1358,8 +1394,7 @@ class ChatStore {
13581394 role : 'assistant' ,
13591395 content : '' ,
13601396 thinking : '' ,
1361- children : [ ] ,
1362- model : currentModelName || undefined
1397+ children : [ ]
13631398 } ,
13641399 userMessageId
13651400 ) ;
0 commit comments