feat(chat): robust thread support (complex)

backend/src/app_mixins/chat.py

@@ -17,6 +17,17 @@ def message_json(msg):
         "threadId": msg.thread_id,
     }
 
+def thread_json(it):
+    return {
+        'id': it.id,
+        'parentId': it.parent_id,
+        'name': it.name,
+        'description': it.description,
+        'firstMessageId': it.first_message_id,
+        'lastMessageId': it.last_message_id,
+        'messageCount': it.message_count,
+    }
+
 async def aggregate_thread(thread_id):
     print(f"aggregate {thread_id}")
     query = db.Message.filter(thread_id=thread_id)
@@ -126,15 +137,7 @@ class ChatMixin:
 
     async def get_chat_threads(self, chat_id):
         threads = await db.Thread.filter(chat_id=chat_id).all()
-        return [{
-            'id': it.id,
-            'parentId': it.parent_id,
-            'name': it.name,
-            'description': it.description,
-            'firstMessageId': it.first_message_id,
-            'lastMessageId': it.last_message_id,
-            'messageCount': it.message_count,
-        } for it in threads]
+        return [thread_json(it) for it in threads]
 
     async def get_thread_messages(self, chat_id, thread_id, limit_message_id=None):
         chat = await self.get_chat(chat_id)
@@ -171,6 +174,16 @@ class ChatMixin:
         for msg in messages:
             self.notify_message_to_user(chat.user_id, msg, True)
 
+
+        # Manually notify thread updates
+        if chat.user_id in self.user_connections:
+            children = await db.Thread.filter(parent_id=new_id).all()
+            packets = [dict(thread_json(it), action='thread-update') for it in children]
+
+            for channel in self.user_connections[chat.user_id]:
+                for packet in packets:
+                    asyncio.create_task(channel.put(packet))
+
         return new_id
 
     async def add_agent_to_chat(self, chat_id, agent_id):

backend/src/inference.py

@@ -154,7 +154,7 @@ async def chat_completion(app,
     agent_run=None
 ):
     # TESTING OVERRIDE
-    return await mock_completion(app, agent_run)
+    # return await mock_completion(app, agent_run)
     db_model = (await app.get_model(model_id)).value
 
     model_engine = db_model['engine']

frontend/src/chat.js

@@ -48,6 +48,19 @@ function selectEarliest (messages) {
   return selected
 }
 
+function messageCopy (it) {
+  return {
+    id: it.id,
+    type: it.kind,
+    format: it.format,
+    subkind: it.subkind,
+    sender: it.sender,
+    content: it.content,
+    threadId: it.threadId,
+    time: it.time,
+  };
+}
+
 export class Thread {
   constructor (obj, session) {
     Object.assign(this, {
@@ -159,6 +172,7 @@ export class ChatSession {
     this._messageNonce = 1;
     this._unsubscribeChat = null;
     this._unsubscribeRuns = null;
+    this._unsubscribeThreads = null;
     this._runListener = null;
     this._rescroll = null;
   }
@@ -211,6 +225,14 @@ export class ChatSession {
 
     this._unsubscribeChat = backend.onChatMessages(chatId, it => this.receiveMessage(it));
     this._unsubscribeRuns = runs.onChatMessage(chatId, event => this.trackRun(event));
+    this._unsubscribeThreads = backend.onAction('thread-update', event => {
+      let thread = this.threadMap[event.id];
+      if (thread) {
+        let data = {...event};
+        delete data.action;
+        thread.update(data)
+      };
+    });
 
     // Load threads
     let {threads} = await backend.sendAction("get-chat-threads", {chatId});
@@ -241,6 +263,8 @@ export class ChatSession {
     this._unsubscribeChat = null;
     if (this._unsubscribeRuns) this._unsubscribeRuns();
     this._unsubscribeRuns = null;
+    if (this._unsubscribeThreads) this._unsubscribeThreads();
+    this._unsubscribeThreads = null;
     if (this._runListener) this._runListener.cancel();
     this._runListener = null;
   }
@@ -435,15 +459,17 @@ export class ChatSession {
 
       if (this.creatingThread.value) {
 
-        // Send the target message to a fork
-        let res = await backend.sendAction("fork-chat-thread", {
-          chatId: this.chatId,
-          threadId: previousThreadId,
-          messageId,
-        });
+        // Send the target message to a fork if it doesn't start its thread
+        if (previousThreadId == messageThreadId) {
+          await backend.sendAction("fork-chat-thread", {
+            chatId: this.chatId,
+            threadId: previousThreadId,
+            messageId,
+          });
+        }
 
-        // Then create a new fork again for the new message
-        res = await backend.sendAction("fork-chat-thread", {
+        // Create a new fork for the new message
+        let res = await backend.sendAction("fork-chat-thread", {
           chatId: this.chatId,
           threadId: previousThreadId,
         });
@@ -598,8 +624,28 @@ export class ChatSession {
     this.currentThreadId.value = threadId;
     this.activeThreadIds = [];
 
-    // Activate and load all threads in the ancestry
     let thread = this.getThread(threadId);
+
+    // If the thread has children, find the most recent leaf
+    let stack = [thread]
+    let lastMessage = null;
+    while (stack.length) {
+      let next = stack.pop();
+
+      let last = next.getLastMessage();
+      if (last && (!lastMessage || last.id > lastMessage.id)) {
+        lastMessage = last;
+      }
+
+      stack = stack.concat(thread.children);
+    }
+
+    // Invariably activate the thread with the most recent message
+    if (lastMessage) {
+      thread = this.getThread(lastMessage.threadId);
+    }
+
+    // Activate and load all threads in the ancestry
     while (true) {
       this.activeThreadIds.push(thread.id);
 
@@ -624,6 +670,8 @@ export class ChatSession {
       thread = this.getThread(thread.parentId);
     }
 
+
+
     this.computeChatEntries();
     this.rescroll();
   }
@@ -732,7 +780,7 @@ export class ChatSession {
           firstMessage ? {id: firstMessage.id, position: 'before'} :
           null;
 
-        if (!attach) return;
+        if (!attach) continue;
 
         headers.push({
           // claim just before id