diff --git a/neo/assistant/api.go b/neo/assistant/api.go
index 6447f82c..41f2b0d4 100644
--- a/neo/assistant/api.go
+++ b/neo/assistant/api.go
@@ -60,7 +60,7 @@ func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string,
 	}
 
 	// Switch to the new assistant if necessary
-	if res.AssistantID != ctx.AssistantID {
+	if res != nil && res.AssistantID != ctx.AssistantID {
 		newAst, err := Get(res.AssistantID)
 		if err != nil {
 			return err
@@ -69,17 +69,17 @@ func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string,
 	}
 
 	// Handle next action
-	if res.Next != nil {
+	if res != nil && res.Next != nil {
 		return res.Next.Execute(c, ctx)
 	}
 
 	// Update options if provided
-	if res.Options != nil {
+	if res != nil && res.Options != nil {
 		options = res.Options
 	}
 
 	// messages
-	if res.Input != nil {
+	if res != nil && res.Input != nil {
 		messages = res.Input
 	}
 
@@ -263,17 +263,20 @@ func (ast *Assistant) streamChat(
 
 				chatMessage.New().
 					Map(map[string]interface{}{
-						"text": value,
-						"done": msg.IsDone,
+						"assistant_id":     ast.ID,
+						"assistant_name":   ast.Name,
+						"assistant_avatar": ast.Avatar,
+						"text":             value,
+						"done":             msg.IsDone,
 					}).
 					Write(c.Writer)
 			}
 
 			// Complete the stream
 			if msg.IsDone {
-				// if value == "" {
-				// 	msg.Write(c.Writer)
-				// }
+				if value == "" {
+					msg.Write(c.Writer)
+				}
 
 				// Call HookDone
 				content.SetStatus(chatMessage.ContentStatusDone)
@@ -300,8 +303,11 @@ func (ast *Assistant) streamChat(
 				} else if value != "" {
 					chatMessage.New().
 						Map(map[string]interface{}{
-							"text": value,
-							"done": true,
+							"assistant_id":     ast.ID,
+							"assistant_name":   ast.Name,
+							"assistant_avatar": ast.Avatar,
+							"text":             value,
+							"done":             true,
 						}).
 						Write(c.Writer)
 				}
diff --git a/neo/message/message.go b/neo/message/message.go
index 6f559bc9..39551102 100644
--- a/neo/message/message.go
+++ b/neo/message/message.go
@@ -16,10 +16,11 @@ import (
 
 // Message the message
 type Message struct {
-	Text            string                 `json:"text,omitempty"`  // text content
-	Type            string                 `json:"type,omitempty"`  // error, text, plan, table, form, page, file, video, audio, image, markdown, json ...
-	Props           map[string]interface{} `json:"props,omitempty"` // props for the types
-	IsDone          bool                   `json:"done,omitempty"`
+	Text            string                 `json:"text,omitempty"`             // text content
+	Type            string                 `json:"type,omitempty"`             // error, text, plan, table, form, page, file, video, audio, image, markdown, json ...
+	Props           map[string]interface{} `json:"props,omitempty"`            // props for the types
+	IsDone          bool                   `json:"done,omitempty"`             // Mark as a done message from neo
+	IsNew           bool                   `json:"is_new,omitempty"`           // Mark as a new message from neo
 	Actions         []Action               `json:"actions,omitempty"`          // Conversation Actions for frontend
 	Attachments     []Attachment           `json:"attachments,omitempty"`      // File attachments
 	Role            string                 `json:"role,omitempty"`             // user, assistant, system ...
@@ -244,6 +245,23 @@ func (m *Message) Map(msg map[string]interface{}) *Message {
 	if done, ok := msg["done"].(bool); ok {
 		m.IsDone = done
 	}
+
+	if isNew, ok := msg["is_new"].(bool); ok {
+		m.IsNew = isNew
+	}
+
+	if assistantID, ok := msg["assistant_id"].(string); ok {
+		m.AssistantID = assistantID
+	}
+
+	if assistantName, ok := msg["assistant_name"].(string); ok {
+		m.AssistantName = assistantName
+	}
+
+	if assistantAvatar, ok := msg["assistant_avatar"].(string); ok {
+		m.AssistantAvatar = assistantAvatar
+	}
+
 	if actions, ok := msg["actions"].([]interface{}); ok {
 		for _, action := range actions {
 			if v, ok := action.(map[string]interface{}); ok {