@liveblocks/core 2.25.0-aiprivatebeta1 → 2.25.0-aiprivatebeta11

package/dist/index.js

@@ -6,7 +6,7 @@ var __export = (target, all) => {
 
 // src/version.ts
 var PKG_NAME = "@liveblocks/core";
-var PKG_VERSION = "2.25.0-aiprivatebeta1";
+var PKG_VERSION = "2.25.0-aiprivatebeta11";
 var PKG_FORMAT = "esm";
 
 // src/dupe-detection.ts
@@ -188,6 +188,9 @@ var warnWithTitle = wrapWithTitle("warn");
 var errorWithTitle = wrapWithTitle("error");
 
 // src/lib/guards.ts
+function isDefined(value) {
+  return value !== null && value !== void 0;
+}
 function isPlainObject(blob) {
   return blob !== null && typeof blob === "object" && Object.prototype.toString.call(blob) === "[object Object]";
 }
@@ -3685,11 +3688,6 @@ function parseAuthToken(rawTokenString) {
 function appendDelta(content, delta) {
   const lastPart = content[content.length - 1];
   switch (delta.type) {
-    case "reasoning":
-    case "text":
-    case "tool-call":
-      content.push(delta);
-      break;
     case "text-delta":
       if (lastPart?.type === "text") {
         lastPart.text += delta.textDelta;
@@ -3700,15 +3698,16 @@ function appendDelta(content, delta) {
     case "reasoning-delta":
       if (lastPart?.type === "reasoning") {
         lastPart.text += delta.textDelta;
-        lastPart.signature ??= delta.signature;
       } else {
         content.push({
           type: "reasoning",
-          text: delta.textDelta ?? "",
-          signature: delta.signature
+          text: delta.textDelta ?? ""
         });
       }
       break;
+    case "tool-invocation":
+      content.push(delta);
+      break;
     default:
       return assertNever(delta, "Unhandled case");
   }
@@ -3716,48 +3715,131 @@ function appendDelta(content, delta) {
 
 // src/ai.ts
 var DEFAULT_REQUEST_TIMEOUT = 4e3;
-var DEFAULT_AI_TIMEOUT = 3e4;
+function defineAiTool() {
+  return (def) => {
+    return def;
+  };
+}
+var KnowledgeStack = class {
+  #_layers;
+  #stack;
+  //                 /                \
+  //      knowledge key               "layer" key
+  //      (random, or optionally      (one entry per mounted component)
+  //       set by user)
+  #_cache;
+  constructor() {
+    this.#_layers = /* @__PURE__ */ new Set();
+    this.#stack = new DefaultMap(
+      () => /* @__PURE__ */ new Map()
+    );
+    this.#_cache = void 0;
+  }
+  // Typically a useId()
+  registerLayer(uniqueLayerId) {
+    const layerKey = uniqueLayerId;
+    if (this.#_layers.has(layerKey))
+      raise(`Layer '${layerKey}' already exists, provide a unique layer id`);
+    this.#_layers.add(layerKey);
+    return layerKey;
+  }
+  deregisterLayer(layerKey) {
+    this.#_layers.delete(layerKey);
+    let deleted = false;
+    for (const [key, knowledge] of this.#stack) {
+      if (knowledge.delete(layerKey)) {
+        deleted = true;
+      }
+      if (knowledge.size === 0)
+        this.#stack.delete(key);
+    }
+    if (deleted) {
+      this.invalidate();
+    }
+  }
+  get() {
+    return this.#_cache ??= this.#recompute();
+  }
+  invalidate() {
+    this.#_cache = void 0;
+  }
+  #recompute() {
+    return Array.from(this.#stack.values()).flatMap(
+      (layer) => (
+        // Return only the last item (returns [] when empty)
+        Array.from(layer.values()).slice(-1).filter(isDefined)
+      )
+    );
+  }
+  updateKnowledge(layerKey, key, data) {
+    if (!this.#_layers.has(layerKey)) raise(`Unknown layer key: ${layerKey}`);
+    this.#stack.getOrCreate(key).set(layerKey, data);
+    this.invalidate();
+  }
+};
 function now() {
   return (/* @__PURE__ */ new Date()).toISOString();
 }
+var kWILDCARD = Symbol("*");
 function createStore_forTools() {
-  const toolsByChatId\u03A3 = new DefaultMap((_chatId) => {
-    return new DefaultMap((_toolName) => {
-      return new Signal(void 0);
+  const toolsByChatId\u03A3 = new DefaultMap(
+    (_chatId) => {
+      return new DefaultMap((_name) => {
+        return new Signal(void 0);
+      });
+    }
+  );
+  const globalOrScopedTool\u03A3 = new DefaultMap((nameAndChat) => {
+    const [name, chatId] = tryParseJson(nameAndChat);
+    return DerivedSignal.from(() => {
+      return (
+        // A tool that's registered and scoped to a specific chat ID...
+        (chatId !== void 0 ? toolsByChatId\u03A3.get(chatId)?.get(name) : void 0)?.get() ?? // ...or a globally registered tool
+        toolsByChatId\u03A3.getOrCreate(kWILDCARD).get(name)?.get()
+      );
     });
   });
-  function getToolDefinition\u03A3(chatId, toolName) {
-    return toolsByChatId\u03A3.getOrCreate(chatId).getOrCreate(toolName);
+  function getTool\u03A3(name, chatId) {
+    const key = JSON.stringify(chatId !== void 0 ? [name, chatId] : [name]);
+    return globalOrScopedTool\u03A3.getOrCreate(key);
   }
-  function addToolDefinition(chatId, name, definition) {
-    toolsByChatId\u03A3.getOrCreate(chatId).getOrCreate(name).set(definition);
+  function registerTool(name, tool, chatId) {
+    if (!tool.execute && !tool.render) {
+      throw new Error(
+        "A tool definition must have an execute() function, a render() function, or both."
+      );
+    }
+    const key = chatId ?? kWILDCARD;
+    toolsByChatId\u03A3.getOrCreate(key).getOrCreate(name).set(tool);
+    return () => unregisterTool(key, name);
   }
-  function removeToolDefinition(chatId, toolName) {
+  function unregisterTool(chatId, name) {
     const tools = toolsByChatId\u03A3.get(chatId);
     if (tools === void 0) return;
-    const tool = tools.get(toolName);
+    const tool = tools.get(name);
     if (tool === void 0) return;
     tool.set(void 0);
   }
-  function getToolsForChat(chatId) {
-    const tools = toolsByChatId\u03A3.get(chatId);
-    if (tools === void 0) return [];
-    return Array.from(tools.entries()).map(([name, tool]) => {
-      if (tool.get() === void 0) return null;
-      return {
-        name,
-        definition: tool.get()
-      };
-    }).filter((tool) => tool !== null);
+  function getToolDescriptions(chatId) {
+    const globalTools\u03A3 = toolsByChatId\u03A3.get(kWILDCARD);
+    const scopedTools\u03A3 = toolsByChatId\u03A3.get(chatId);
+    return Array.from([
+      ...globalTools\u03A3?.entries() ?? [],
+      ...scopedTools\u03A3?.entries() ?? []
+    ]).flatMap(([name, tool\u03A3]) => {
+      const tool = tool\u03A3.get();
+      return tool ? [{ name, description: tool.description, parameters: tool.parameters }] : [];
+    });
   }
   return {
-    getToolCallByName\u03A3: getToolDefinition\u03A3,
-    getToolsForChat,
-    addToolDefinition,
-    removeToolDefinition
+    getToolDescriptions,
+    getTool\u03A3,
+    registerTool
   };
 }
-function createStore_forChatMessages() {
+function createStore_forChatMessages(toolsStore, setToolResult) {
+  const autoExecutableMessages = /* @__PURE__ */ new Set();
+  const seenToolCallIds = /* @__PURE__ */ new Set();
   const messagePoolByChatId\u03A3 = new DefaultMap(
     (_chatId) => new MutableSignal(
       new TreePool(
@@ -3767,7 +3849,7 @@ function createStore_forChatMessages() {
       )
     )
   );
-  const pendingMessages\u03A3 = new MutableSignal(
+  const generatingMessages\u03A3 = new MutableSignal(
     /* @__PURE__ */ new Map()
   );
   function createOptimistically(chatId, role, parentId, third) {
@@ -3781,7 +3863,8 @@ function createStore_forChatMessages() {
         role,
         parentId,
         createdAt,
-        content
+        content,
+        _optimistic: true
       });
     } else {
       upsert({
@@ -3790,8 +3873,9 @@ function createStore_forChatMessages() {
         role,
         parentId,
         createdAt,
-        status: "pending",
-        contentSoFar: []
+        status: "generating",
+        contentSoFar: [],
+        _optimistic: true
       });
     }
     return id;
@@ -3808,7 +3892,7 @@ function createStore_forChatMessages() {
     if (!chatMsgs\u03A3) return;
     const existing = chatMsgs\u03A3.get().get(messageId);
     if (!existing || existing.deletedAt) return;
-    if (existing.role === "assistant" && (existing.status === "pending" || existing.status === "failed")) {
+    if (existing.role === "assistant" && existing.status !== "completed") {
       upsert({ ...existing, deletedAt: now(), contentSoFar: [] });
     } else {
       upsert({ ...existing, deletedAt: now(), content: [] });
@@ -3823,19 +3907,59 @@ function createStore_forChatMessages() {
     batch(() => {
       const chatMsgs\u03A3 = messagePoolByChatId\u03A3.getOrCreate(message.chatId);
       chatMsgs\u03A3.mutate((pool) => pool.upsert(message));
-      if (message.role === "assistant" && message.status === "pending") {
-        pendingMessages\u03A3.mutate((lut) => {
+      if (message.role === "assistant" && message.status === "generating") {
+        generatingMessages\u03A3.mutate((lut) => {
           lut.set(message.id, structuredClone(message));
         });
       } else {
-        pendingMessages\u03A3.mutate((lut) => {
+        generatingMessages\u03A3.mutate((lut) => {
           lut.delete(message.id);
         });
       }
+      if (message.role === "assistant" && message.status === "awaiting-tool") {
+        for (const toolCall of message.contentSoFar.filter(
+          (part) => part.type === "tool-invocation" && part.status === "executing"
+        )) {
+          if (seenToolCallIds.has(toolCall.toolCallId)) {
+            continue;
+          }
+          seenToolCallIds.add(toolCall.toolCallId);
+          const toolDef = toolsStore.getTool\u03A3(toolCall.toolName, message.chatId).get();
+          const respondSync = (result) => {
+            setToolResult(
+              message.chatId,
+              message.id,
+              toolCall.toolCallId,
+              result
+              // TODO Pass in AiGenerationOptions here, or make the backend use the same options
+            ).catch((err) => {
+              error2(
+                `Error trying to respond to tool-call: ${String(err)} (in respond())`
+              );
+            });
+          };
+          const executeFn = toolDef?.execute;
+          if (executeFn && autoExecutableMessages.has(message.id)) {
+            (async () => {
+              const result = await executeFn(toolCall.args, {
+                toolName: toolCall.toolName,
+                toolCallId: toolCall.toolCallId
+              });
+              respondSync(result);
+            })().catch((err) => {
+              error2(
+                `Error trying to respond to tool-call: ${String(err)} (in execute())`
+              );
+            });
+          }
+        }
+      } else {
+        autoExecutableMessages.delete(message.id);
+      }
     });
   }
   function addDelta(messageId, delta) {
-    pendingMessages\u03A3.mutate((lut) => {
+    generatingMessages\u03A3.mutate((lut) => {
       const message = lut.get(messageId);
       if (message === void 0) return false;
       appendDelta(message.contentSoFar, delta);
@@ -3843,10 +3967,10 @@ function createStore_forChatMessages() {
       return true;
     });
   }
-  function* iterPendingMessages() {
+  function* iterGeneratingMessages() {
     for (const chatMsgs\u03A3 of messagePoolByChatId\u03A3.values()) {
       for (const m of chatMsgs\u03A3.get()) {
-        if (m.role === "assistant" && m.status === "pending") {
+        if (m.role === "assistant" && m.status === "generating" && !m._optimistic) {
           yield m;
         }
       }
@@ -3854,9 +3978,18 @@ function createStore_forChatMessages() {
   }
   function failAllPending() {
     batch(() => {
-      pendingMessages\u03A3.mutate((lut) => lut.clear());
+      generatingMessages\u03A3.mutate((lut) => {
+        let deleted = false;
+        for (const [k, v] of lut) {
+          if (!v._optimistic) {
+            lut.delete(k);
+            deleted = true;
+          }
+        }
+        return deleted;
+      });
       upsertMany(
-        Array.from(iterPendingMessages()).map(
+        Array.from(iterGeneratingMessages()).map(
           (message) => ({
             ...message,
             status: "failed",
@@ -3891,11 +4024,20 @@ function createStore_forChatMessages() {
     }
     function selectSpine(leaf) {
       const spine = [];
+      let lastVisitedMessage = null;
       for (const message2 of pool.walkUp(leaf.id)) {
         const prev = first(pool.walkLeft(message2.id, isAlive))?.id ?? null;
         const next = first(pool.walkRight(message2.id, isAlive))?.id ?? null;
         if (!message2.deletedAt || prev || next) {
-          spine.push({ ...message2, prev, next });
+          const node = {
+            ...message2,
+            navigation: { parent: null, prev, next }
+          };
+          if (lastVisitedMessage !== null) {
+            lastVisitedMessage.navigation.parent = node.id;
+          }
+          lastVisitedMessage = node;
+          spine.push(node);
         }
       }
       return spine.reverse();
@@ -3921,18 +4063,6 @@ function createStore_forChatMessages() {
     }
     return fallback();
   }
-  function getLatestUserMessageAncestor(chatId, messageId) {
-    const pool = messagePoolByChatId\u03A3.getOrCreate(chatId).get();
-    const message = pool.get(messageId);
-    if (!message) return null;
-    if (message.role === "user") return message.id;
-    for (const m of pool.walkUp(message.id)) {
-      if (m.role === "user" && !m.deletedAt) {
-        return m.id;
-      }
-    }
-    return null;
-  }
   const immutableMessagesByBranch = new DefaultMap((chatId) => {
     return new DefaultMap((branchId) => {
       const messages\u03A3 = DerivedSignal.from(() => {
@@ -3940,16 +4070,16 @@ function createStore_forChatMessages() {
         return selectBranch(pool, branchId);
       }, shallow2);
       return DerivedSignal.from(() => {
-        const pendingMessages = pendingMessages\u03A3.get();
+        const generatingMessages = generatingMessages\u03A3.get();
         return messages\u03A3.get().map((message) => {
-          if (message.role !== "assistant" || message.status !== "pending") {
+          if (message.role !== "assistant" || message.status !== "generating") {
             return message;
           }
-          const pendingMessage = pendingMessages.get(message.id);
-          if (pendingMessage === void 0) return message;
+          const generatingMessage = generatingMessages.get(message.id);
+          if (generatingMessage === void 0) return message;
           return {
             ...message,
-            contentSoFar: pendingMessage.contentSoFar
+            contentSoFar: generatingMessage.contentSoFar
           };
         });
       }, shallow);
@@ -3958,21 +4088,10 @@ function createStore_forChatMessages() {
   function getChatMessagesForBranch\u03A3(chatId, branch) {
     return immutableMessagesByBranch.getOrCreate(chatId).getOrCreate(branch || null);
   }
-  const messagesByChatId\u03A3 = new DefaultMap((chatId) => {
-    return DerivedSignal.from(() => {
-      const pool = messagePoolByChatId\u03A3.getOrCreate(chatId).get();
-      return Array.from(pool.sorted);
-    });
-  });
-  function getMessagesForChat\u03A3(chatId) {
-    return messagesByChatId\u03A3.getOrCreate(chatId);
-  }
   return {
     // Readers
     getMessageById,
     getChatMessagesForBranch\u03A3,
-    getMessagesForChat\u03A3,
-    getLatestUserMessageAncestor,
     // Mutations
     createOptimistically,
     upsert,
@@ -3980,20 +4099,23 @@ function createStore_forChatMessages() {
     remove,
     removeByChatId,
     addDelta,
-    failAllPending
+    failAllPending,
+    allowAutoExecuteToolCall(messageId) {
+      autoExecutableMessages.add(messageId);
+    }
   };
 }
 function createStore_forUserAiChats() {
-  const mutable\u03A3 = new MutableSignal(
+  const allChatsInclDeleted\u03A3 = new MutableSignal(
     SortedList.with((x, y) => y.createdAt < x.createdAt)
   );
-  const chats\u03A3 = DerivedSignal.from(
-    () => Array.from(mutable\u03A3.get()).filter((c) => !c.ephemeral && !c.deletedAt)
+  const nonDeletedChats\u03A3 = DerivedSignal.from(
+    () => Array.from(allChatsInclDeleted\u03A3.get()).filter((c) => !c.deletedAt)
   );
   function upsertMany(chats) {
-    mutable\u03A3.mutate((list) => {
+    allChatsInclDeleted\u03A3.mutate((list) => {
       for (const chat of chats) {
-        remove(chat.id);
+        list.removeBy((c) => c.id === chat.id, 1);
         list.add(chat);
       }
     });
@@ -4001,15 +4123,26 @@ function createStore_forUserAiChats() {
   function upsert(chat) {
     upsertMany([chat]);
   }
-  function remove(chatId) {
-    mutable\u03A3.mutate((list) => list.removeBy((c) => c.id === chatId, 1));
+  function markDeleted(chatId) {
+    allChatsInclDeleted\u03A3.mutate((list) => {
+      const chat = list.find((c) => c.id === chatId);
+      if (!chat) return false;
+      upsert({ ...chat, deletedAt: now() });
+      return void 0;
+    });
+  }
+  function getChatById(chatId) {
+    return Array.from(allChatsInclDeleted\u03A3.get()).find(
+      (chat) => chat.id === chatId
+    );
   }
   return {
-    chats\u03A3,
+    chats\u03A3: nonDeletedChats\u03A3,
+    getChatById,
     // Mutations
     upsert,
     upsertMany,
-    remove
+    markDeleted
   };
 }
 function createAi(config) {
@@ -4019,10 +4152,9 @@ function createAi(config) {
     false
     // AI doesn't have actors (yet, but it will)
   );
-  const clientId = nanoid(7);
   const chatsStore = createStore_forUserAiChats();
-  const messagesStore = createStore_forChatMessages();
   const toolsStore = createStore_forTools();
+  const messagesStore = createStore_forChatMessages(toolsStore, setToolResult);
   const context = {
     staticSessionInfoSig: new Signal(null),
     dynamicSessionInfoSig: new Signal(null),
@@ -4030,11 +4162,10 @@ function createAi(config) {
     chatsStore,
     messagesStore,
     toolsStore,
-    contextByChatId: /* @__PURE__ */ new Map()
+    knowledge: new KnowledgeStack()
   };
   let lastTokenKey;
-  function onStatusDidChange(newStatus) {
-    warn("onStatusDidChange", newStatus);
+  function onStatusDidChange(_newStatus) {
     const authValue = managedSocket.authValue;
     if (authValue !== null) {
       const tokenKey = getBearerTokenFromAuthValue(authValue);
@@ -4070,7 +4201,6 @@ function createAi(config) {
     }
   }
   function onDidConnect() {
-    warn("onDidConnect");
   }
   function onDidDisconnect() {
     warn("onDidDisconnect");
@@ -4112,7 +4242,7 @@ function createAi(config) {
               context.messagesStore.remove(m.chatId, m.id);
             }
             for (const chatId of msg["-chats"] ?? []) {
-              context.chatsStore.remove(chatId);
+              context.chatsStore.markDeleted(chatId);
               context.messagesStore.removeByChatId(chatId);
             }
             for (const chatId of msg.clear ?? []) {
@@ -4134,34 +4264,36 @@ function createAi(config) {
         case "get-chats":
           context.chatsStore.upsertMany(msg.chats);
           break;
-        case "create-chat":
+        case "get-or-create-chat":
           context.chatsStore.upsert(msg.chat);
           break;
         case "delete-chat":
-          context.chatsStore.remove(msg.chatId);
+          context.chatsStore.markDeleted(msg.chatId);
           context.messagesStore.removeByChatId(msg.chatId);
           break;
         case "get-message-tree":
           context.chatsStore.upsert(msg.chat);
           context.messagesStore.upsertMany(msg.messages);
           break;
-        case "add-user-message":
-          context.messagesStore.upsert(msg.message);
-          break;
         case "delete-message":
           context.messagesStore.remove(msg.chatId, msg.messageId);
           break;
         case "clear-chat":
           context.messagesStore.removeByChatId(msg.chatId);
           break;
-        case "ask-ai":
-          if (msg.message) {
-            context.messagesStore.upsert(msg.message);
-          } else {
+        case "ask-in-chat":
+          if (msg.sourceMessage) {
+            context.messagesStore.upsert(msg.sourceMessage);
           }
+          context.messagesStore.upsert(msg.targetMessage);
           break;
         case "abort-ai":
           break;
+        case "set-tool-result":
+          if (msg.ok) {
+            context.messagesStore.upsert(msg.message);
+          }
+          break;
         default:
           return assertNever(msg, "Unhandled case");
       }
@@ -4212,13 +4344,11 @@ function createAi(config) {
       cursor: options.cursor
     });
   }
-  function createChat(id, name, options) {
+  function getOrCreateChat(id, options) {
     return sendClientMsgWithResponse({
-      cmd: "create-chat",
+      cmd: "get-or-create-chat",
       id,
-      name,
-      ephemeral: options?.ephemeral ?? false,
-      metadata: options?.metadata ?? {}
+      options
     });
   }
   function getMessageTree(chatId) {
@@ -4227,147 +4357,89 @@ function createAi(config) {
       chatId
     });
   }
-  function registerChatContext(chatId, data) {
-    const chatContext = context.contextByChatId.get(chatId);
-    if (chatContext === void 0) {
-      context.contextByChatId.set(chatId, /* @__PURE__ */ new Set([data]));
-    } else {
-      chatContext.add(data);
-    }
-    return () => {
-      const chatContext2 = context.contextByChatId.get(chatId);
-      if (chatContext2 !== void 0) {
-        chatContext2.delete(data);
-        if (chatContext2.size === 0) {
-          context.contextByChatId.delete(chatId);
-        }
-      }
-    };
+  function registerKnowledgeLayer(uniqueLayerId) {
+    return context.knowledge.registerLayer(uniqueLayerId);
   }
-  function ask(chatId, messageId, options) {
-    const targetMessageId = context.messagesStore.createOptimistically(
-      chatId,
-      "assistant",
-      messageId
-    );
-    const copilotId = options?.copilotId;
-    const stream = options?.stream ?? false;
-    const timeout = options?.timeout ?? DEFAULT_AI_TIMEOUT;
-    const chatContext = context.contextByChatId.get(chatId);
-    return sendClientMsgWithResponse({
-      cmd: "ask-ai",
+  function deregisterKnowledgeLayer(layerKey) {
+    context.knowledge.deregisterLayer(layerKey);
+  }
+  function updateKnowledge(layerKey, data, key = nanoid()) {
+    context.knowledge.updateKnowledge(layerKey, key, data);
+  }
+  async function setToolResult(chatId, messageId, toolCallId, result, options) {
+    const knowledge = context.knowledge.get();
+    const tools = context.toolsStore.getToolDescriptions(chatId);
+    const resp = await sendClientMsgWithResponse({
+      cmd: "set-tool-result",
       chatId,
-      sourceMessageId: messageId,
-      targetMessageId,
-      copilotId,
-      clientId,
-      stream,
-      tools: context.toolsStore.getToolsForChat(chatId).map((tool) => ({
-        name: tool.name,
-        description: tool.definition.description,
-        parameters: tool.definition.parameters
-      })),
-      timeout,
-      context: chatContext ? Array.from(chatContext.values()) : void 0
+      messageId,
+      toolCallId,
+      result,
+      generationOptions: {
+        copilotId: options?.copilotId,
+        stream: options?.stream,
+        timeout: options?.timeout,
+        // Knowledge and tools aren't coming from the options, but retrieved
+        // from the global context
+        knowledge: knowledge.length > 0 ? knowledge : void 0,
+        tools: tools.length > 0 ? tools : void 0
+      }
     });
+    if (resp.ok) {
+      messagesStore.allowAutoExecuteToolCall(resp.message.id);
+    }
+    return resp;
   }
   return Object.defineProperty(
     {
       [kInternal]: {
-        debugContext: () => context
+        context
       },
       connect: () => managedSocket.connect(),
       reconnect: () => managedSocket.reconnect(),
       disconnect: () => managedSocket.disconnect(),
       getChats,
-      createChat,
+      getOrCreateChat,
       deleteChat: (chatId) => {
-        return sendClientMsgWithResponse({
-          cmd: "delete-chat",
-          chatId
-        });
+        return sendClientMsgWithResponse({ cmd: "delete-chat", chatId });
       },
       getMessageTree,
       deleteMessage: (chatId, messageId) => sendClientMsgWithResponse({ cmd: "delete-message", chatId, messageId }),
       clearChat: (chatId) => sendClientMsgWithResponse({ cmd: "clear-chat", chatId }),
-      addUserMessage: (chatId, parentMessageId, message) => {
-        const content = [{ type: "text", text: message }];
-        const newMessageId = context.messagesStore.createOptimistically(
-          chatId,
-          "user",
-          parentMessageId,
-          content
-        );
-        return sendClientMsgWithResponse({
-          cmd: "add-user-message",
-          id: newMessageId,
+      askUserMessageInChat: async (chatId, userMessage, targetMessageId, options) => {
+        const knowledge = context.knowledge.get();
+        const tools = context.toolsStore.getToolDescriptions(chatId);
+        const resp = await sendClientMsgWithResponse({
+          cmd: "ask-in-chat",
           chatId,
-          parentMessageId,
-          content
-        });
-      },
-      ask,
-      regenerateMessage: (chatId, messageId, options) => {
-        const parentUserMessageId = context.messagesStore.getLatestUserMessageAncestor(chatId, messageId);
-        if (parentUserMessageId === null) {
-          throw new Error(
-            `Unable to find user message ancestor for messageId: ${messageId}`
-          );
-        }
-        return ask(chatId, parentUserMessageId, options);
-      },
-      addUserMessageAndAsk: async (chatId, parentMessageId, message, options) => {
-        const content = [{ type: "text", text: message }];
-        const newMessageId = context.messagesStore.createOptimistically(
-          chatId,
-          "user",
-          parentMessageId,
-          content
-        );
-        const targetMessageId = context.messagesStore.createOptimistically(
-          chatId,
-          "assistant",
-          newMessageId
-        );
-        await sendClientMsgWithResponse({
-          cmd: "add-user-message",
-          id: newMessageId,
-          chatId,
-          parentMessageId,
-          content
-        });
-        const copilotId = options?.copilotId;
-        const stream = options?.stream ?? false;
-        const timeout = options?.timeout ?? DEFAULT_AI_TIMEOUT;
-        const chatContext = context.contextByChatId.get(chatId);
-        return sendClientMsgWithResponse({
-          cmd: "ask-ai",
-          chatId,
-          sourceMessageId: newMessageId,
+          sourceMessage: userMessage,
           targetMessageId,
-          copilotId,
-          clientId,
-          stream,
-          tools: context.toolsStore.getToolsForChat(chatId).map((tool) => ({
-            name: tool.name,
-            description: tool.definition.description,
-            parameters: tool.definition.parameters
-          })),
-          timeout,
-          context: chatContext ? Array.from(chatContext.values()) : void 0
+          generationOptions: {
+            copilotId: options?.copilotId,
+            stream: options?.stream,
+            timeout: options?.timeout,
+            // Knowledge and tools aren't coming from the options, but retrieved
+            // from the global context
+            knowledge: knowledge.length > 0 ? knowledge : void 0,
+            tools: tools.length > 0 ? tools : void 0
+          }
         });
+        messagesStore.allowAutoExecuteToolCall(resp.targetMessage.id);
+        return resp;
       },
       abort: (messageId) => sendClientMsgWithResponse({ cmd: "abort-ai", messageId }),
+      setToolResult,
       getStatus: () => managedSocket.getStatus(),
       signals: {
         chats\u03A3: context.chatsStore.chats\u03A3,
         getChatMessagesForBranch\u03A3: context.messagesStore.getChatMessagesForBranch\u03A3,
-        getToolDefinition\u03A3: context.toolsStore.getToolCallByName\u03A3,
-        getMessagesForChat\u03A3: context.messagesStore.getMessagesForChat\u03A3
+        getTool\u03A3: context.toolsStore.getTool\u03A3
       },
-      registerChatContext,
-      registerChatTool: context.toolsStore.addToolDefinition,
-      unregisterChatTool: context.toolsStore.removeToolDefinition
+      getChatById: context.chatsStore.getChatById,
+      registerKnowledgeLayer,
+      deregisterKnowledgeLayer,
+      updateKnowledge,
+      registerTool: context.toolsStore.registerTool
     },
     kInternal,
     { enumerable: false }
@@ -10531,6 +10603,7 @@ export {
   createManagedPool,
   createNotificationSettings,
   createThreadId,
+  defineAiTool,
   deprecate,
   deprecateIf,
   detectDupes,